1// SPDX-License-Identifier: MIT
2
3package guest
4
5import (
6 "context"
7 "errors"
8 "fmt"
9 "io"
10 "net"
11 "sync"
12 "time"
13)
14
15// The cache service of the runner never listens on a network. In the sandbox it
16// is reached on the guest's own loopback (CacheListenHost, on the port the
17// host gives in JobSpec.cache_port), where a Relay accepts the TCP connections
18// of the actions and pipes each of them to a new stream on the host↔guest
19// channel (a vsock port for a VM, a Unix socket for a container); the host
20// serves HTTP on the other end of that stream.
21const (
22 // CacheListenHost is the loopback address of the relay.
23 CacheListenHost = "127.0.0.1"
24 // DefaultRelayMax bounds the connections a Relay pipes at once. The host
25 // bounds them too (cachesvc.DefaultConnsPerListener): a connection over
26 // either limit is closed at once and the client retries.
27 DefaultRelayMax = 32
28 // relayDialTimeout bounds one connection to the host: a stream either opens
29 // at once or fails.
30 relayDialTimeout = 5 * time.Second
31 // DefaultHalfCloseGrace is how long the second direction of a connection may
32 // stay idle once the first has ended. A half-close is passed on, but whether
33 // every channel carries it is not something the relay can rely on (it is
34 // untested over Firecracker's vsock proxy): without a limit, a client that
35 // closed its connection could hold a slot until the host's idle timeout. A
36 // transfer that makes progress is never idle, so only a stalled half is cut.
37 DefaultHalfCloseGrace = 30 * time.Second
38)
39
40// DialFunc opens one new stream to the host.
41type DialFunc func(ctx context.Context) (io.ReadWriteCloser, error)
42
43// Relay accepts TCP connections on a loopback address and pipes each one to a
44// new stream from its DialFunc, in both directions, until Close.
45type Relay struct {
46 ln net.Listener
47 dial DialFunc
48 slot chan struct{}
49 grace time.Duration
50
51 ctx context.Context
52 cancel context.CancelFunc
53 wg sync.WaitGroup
54
55 mu sync.Mutex
56 active map[io.Closer]struct{}
57}
58
59// StartRelay listens on addr ("127.0.0.1:9089") and relays every connection to
60// a stream made by dial; at most max connections are open at once (<= 0 means
61// DefaultRelayMax). The relay ends with Close or when ctx ends.
62func StartRelay(ctx context.Context, addr string, dial DialFunc, max int) (*Relay, error) {
63 return newRelay(ctx, addr, dial, max, DefaultHalfCloseGrace)
64}
65
66func newRelay(ctx context.Context, addr string, dial DialFunc, max int, grace time.Duration) (*Relay, error) {
67 if max <= 0 {
68 max = DefaultRelayMax
69 }
70 ln, err := net.Listen("tcp", addr)
71 if err != nil {
72 return nil, fmt.Errorf("cache relay: %w", err)
73 }
74 rctx, cancel := context.WithCancel(ctx)
75 r := &Relay{ln: ln, dial: dial, slot: make(chan struct{}, max), grace: grace, ctx: rctx, cancel: cancel, active: map[io.Closer]struct{}{}}
76 r.wg.Add(1)
77 go r.accept()
78 context.AfterFunc(rctx, func() { _ = r.Close() })
79 return r, nil
80}
81
82// Addr is the address the relay listens on.
83func (r *Relay) Addr() net.Addr { return r.ln.Addr() }
84
85func (r *Relay) accept() {
86 defer r.wg.Done()
87 for {
88 c, err := r.ln.Accept()
89 if err != nil {
90 if r.ctx.Err() != nil || errors.Is(err, net.ErrClosed) {
91 return
92 }
93 // A transient accept error (descriptors, aborted connection): go on.
94 select {
95 case <-time.After(10 * time.Millisecond):
96 continue
97 case <-r.ctx.Done():
98 return
99 }
100 }
101 select {
102 case r.slot <- struct{}{}:
103 default:
104 _ = c.Close() // over the cap: never queued
105 continue
106 }
107 r.wg.Add(1)
108 go func() {
109 defer r.wg.Done()
110 defer func() { <-r.slot }()
111 r.pipe(c)
112 }()
113 }
114}
115
116// track registers a connection to be closed by Close; it reports false (and
117// closes c) when the relay is already closing.
118func (r *Relay) track(c io.Closer) bool {
119 r.mu.Lock()
120 defer r.mu.Unlock()
121 if r.ctx.Err() != nil {
122 _ = c.Close()
123 return false
124 }
125 r.active[c] = struct{}{}
126 return true
127}
128
129func (r *Relay) untrack(c io.Closer) {
130 r.mu.Lock()
131 delete(r.active, c)
132 r.mu.Unlock()
133 _ = c.Close()
134}
135
136func (r *Relay) pipe(client net.Conn) {
137 if !r.track(client) {
138 return
139 }
140 defer r.untrack(client)
141 dctx, cancel := context.WithTimeout(r.ctx, relayDialTimeout)
142 host, err := r.dial(dctx)
143 cancel()
144 if err != nil {
145 return
146 }
147 if !r.track(host) {
148 return
149 }
150 defer r.untrack(host)
151 Pipe(client, host, r.grace)
152}
153
154// closeWriter is a stream that can shut down its write side only.
155type closeWriter interface{ CloseWrite() error }
156
157// Pipe copies a to b and b to a until both directions end. The end of one
158// direction is passed on as a half-close of the other side's write end, so a
159// request sent in full and answered afterwards works; a stream that cannot
160// half-close is closed once its peer is done. Once one direction has ended, the
161// other may stay idle for grace (0: no limit) before both are closed. Pipe
162// closes a and b only in that case; the caller closes them otherwise.
163func Pipe(a, b io.ReadWriteCloser, grace time.Duration) {
164 var (
165 wg sync.WaitGroup
166 mu sync.Mutex
167 timer *time.Timer
168 done int
169 )
170 touch := func() {
171 mu.Lock()
172 if timer != nil {
173 timer.Reset(grace)
174 }
175 mu.Unlock()
176 }
177 ended := func() {
178 mu.Lock()
179 defer mu.Unlock()
180 done++
181 switch {
182 case done == 1 && grace > 0:
183 timer = time.AfterFunc(grace, func() { _ = a.Close(); _ = b.Close() })
184 case done == 2 && timer != nil:
185 timer.Stop()
186 }
187 }
188 cp := func(dst, src io.ReadWriteCloser) {
189 defer wg.Done()
190 _, _ = io.Copy(activityWriter{dst, touch}, src)
191 if cw, ok := dst.(closeWriter); ok {
192 _ = cw.CloseWrite()
193 } else {
194 _ = dst.Close()
195 }
196 ended()
197 }
198 wg.Add(2)
199 go cp(a, b)
200 go cp(b, a)
201 wg.Wait()
202}
203
204// activityWriter calls touch after every write that goes through.
205type activityWriter struct {
206 io.Writer
207 touch func()
208}
209
210func (w activityWriter) Write(p []byte) (int, error) {
211 n, err := w.Writer.Write(p)
212 if n > 0 {
213 w.touch()
214 }
215 return n, err
216}
217
218// Close stops listening, closes every open connection and waits for the pipes
219// to end. It is idempotent.
220func (r *Relay) Close() error {
221 r.cancel()
222 err := r.ln.Close()
223 r.mu.Lock()
224 for c := range r.active {
225 _ = c.Close()
226 }
227 r.mu.Unlock()
228 r.wg.Wait()
229 if errors.Is(err, net.ErrClosed) {
230 return nil
231 }
232 return err
233}