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