diff --git a/.golangci.yml b/.golangci.yml index 1eadd9508..ef776a732 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -39,7 +39,6 @@ linters-settings: desc: The io/ioutil package has been deprecated. forbidigo: forbid: - - '^fmt\.Errorf(# use errors\.Errorf instead)?$' - '^logrus\.(Trace|Debug|Info|Warn|Warning|Error|Fatal)(f|ln)?(# use bklog\.G or bklog\.L instead of logrus directly)?$' importas: alias: diff --git a/api/services/registry/copy.go b/api/services/registry/copy.go new file mode 100644 index 000000000..6a58ee5f9 --- /dev/null +++ b/api/services/registry/copy.go @@ -0,0 +1,132 @@ +package earthly_registry_v1 //nolint:revive + +import ( + "context" + "errors" + "fmt" + "io" + + "golang.org/x/sync/errgroup" +) + +// copyBufferSize is the size of the buffer each direction reads into. 32KiB is +// io.Copy's own default, and what session/sshforward.Copy and +// session/socketforward use for the same job; it is comfortably under gRPC's +// 4MiB default maximum message size, so a full buffer never needs splitting +// across messages. Nothing depends on the two ends of the tunnel choosing the +// same value -- that they happened to agree is what kept StreamRW's leftover +// handling from ever being exercised -- so this is a throughput knob and +// nothing more. +const copyBufferSize = 32 * 1024 + +// Stream is the part of a gRPC bidirectional stream the tunnel needs. Both +// ends of Registry.Proxy satisfy it, so the same copy runs on the daemon and +// on the client. +type Stream interface { + SendMsg(m any) error + RecvMsg(m any) error +} + +// Copy joins a connection to a gRPC stream in both directions and returns once +// both are done. The bytes are opaque: nothing here knows where one HTTP +// request or response ends, because nothing needs to. Each direction ends when +// its source says so -- io.EOF from the connection, or a peer that closed its +// send direction -- and that end is passed on as a half-close, so the other +// side can finish what it still owes before the whole conversation is torn +// down. +// +// This mirrors session/sshforward.Copy, which has carried forwarded agent +// sockets for years; see that file for the same shape with commentary. +func Copy(ctx context.Context, conn io.ReadWriteCloser, stream Stream, closeStream func() error) error { + defer conn.Close() + + eg, ctx := errgroup.WithContext(ctx) + + // Peer to connection. + eg.Go(func() error { + msg := &ByteMessage{} + for { + if err := stream.RecvMsg(msg); err != nil { + if errors.Is(err, io.EOF) { + // The peer has finished sending. It is still reading, so + // close only this direction and leave the response to + // come back. + if closeWriter, ok := conn.(interface{ CloseWrite() error }); ok { + // Best effort: the read half stays open either way. + closeWriter.CloseWrite() + } else { + conn.Close() + } + return nil + } + conn.Close() + return fmt.Errorf("receive from stream: %w", err) + } + + select { + case <-ctx.Done(): + conn.Close() + return context.Cause(ctx) + default: + } + + if _, err := conn.Write(msg.GetData()); err != nil { + conn.Close() + return fmt.Errorf("write to connection: %w", err) + } + + msg.Data = msg.Data[:0] + } + }) + + // Connection to peer. + eg.Go(func() error { + buf := make([]byte, copyBufferSize) + + send := func(n int) error { + if n == 0 { + return nil + } + if err := stream.SendMsg(&ByteMessage{Data: buf[:n]}); err != nil { + return fmt.Errorf("send to stream: %w", err) + } + return nil + } + + for { + n, err := conn.Read(buf) + switch { + case errors.Is(err, io.EOF): + // Everything the connection had to say has been said. A final + // read is allowed to hand back bytes alongside the EOF, so + // they still go out before the stream is closed. + if err := send(n); err != nil { + return err + } + if closeStream != nil { + if err := closeStream(); err != nil { + return fmt.Errorf("close stream: %w", err) + } + } + return nil + case err != nil: + // A read error on the connection is terminal. Whatever is in + // the buffer belongs to a response that will never be + // completed, so there is nothing worth forwarding. + return fmt.Errorf("read from connection: %w", err) + } + + if err := send(n); err != nil { + return err + } + + select { + case <-ctx.Done(): + return context.Cause(ctx) + default: + } + } + }) + + return eg.Wait() +} diff --git a/api/services/registry/proxy_test.go b/api/services/registry/proxy_test.go new file mode 100644 index 000000000..1aeff1938 --- /dev/null +++ b/api/services/registry/proxy_test.go @@ -0,0 +1,444 @@ +package earthly_registry_v1 //nolint:revive + +import ( + "bytes" + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "strconv" + "sync" + "testing" + "time" + + "google.golang.org/grpc" +) + +var errUnexpectedMessage = errors.New("stream carried a message that was not a ByteMessage") + +// The proxy carries one HTTP conversation between a local docker daemon and +// buildkitd's embedded registry over a gRPC stream, as opaque bytes. These +// tests stand a real HTTP server in for the embedded registry, drive it with a +// real net/http client, and put the production Server.Proxy between the two. +// Nothing here parses HTTP: whether the bytes arrive intact is the whole +// question, so both ends are left to net/http to frame. + +// tunnel returns an http.Client whose connections are proxied to reg through +// Server.Proxy, and a count of how many connections it dialled -- reuse of a +// single connection is a property worth asserting, not assuming. +func tunnel(t *testing.T, reg *httptest.Server) (*http.Client, func() int) { + t.Helper() + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + srv := NewServer(reg.Listener.Addr().String()) + + var ( + mu sync.Mutex + dials int + ) + + tr := &http.Transport{ + DialContext: func(_ context.Context, _, _ string) (net.Conn, error) { + mu.Lock() + dials++ + mu.Unlock() + + near, far, err := tcpPair(t) + if err != nil { + return nil, err + } + + clientStream, serverStream := newStreamPair(ctx) + + go func() { + defer serverStream.closeSend() + _ = srv.Proxy(serverStream) + }() + + go func() { + defer far.Close() + _ = pump(far, clientStream) + }() + + return near, nil + }, + } + t.Cleanup(tr.CloseIdleConnections) + + // A tunnel that mishandles termination strands a request rather than + // failing it, so bound every request: a hung pull is a failure too. + return &http.Client{Transport: tr, Timeout: 20 * time.Second}, func() int { + mu.Lock() + defer mu.Unlock() + return dials + } +} + +// pump is the client half of the tunnel: the daemon-side counterpart to +// Server.Proxy, written out here rather than borrowed from production code so +// that the tests exercise the server against an independently correct peer. +// Termination is by half-close in both directions, never by a timer. +func pump(conn *net.TCPConn, stream *clientStream) error { + errs := make(chan error, 2) + + go func() { + buf := make([]byte, 32*1024) + for { + n, err := conn.Read(buf) + if n > 0 { + if serr := stream.SendMsg(&ByteMessage{Data: buf[:n]}); serr != nil { + errs <- serr + return + } + } + if err != nil { + // The daemon has finished its request. Tell the far side, but + // keep reading the response. + stream.closeSend() + if err == io.EOF { + err = nil + } + errs <- err + return + } + } + }() + + go func() { + for { + msg := &ByteMessage{} + err := stream.RecvMsg(msg) + if err != nil { + if err == io.EOF { + // The response is complete: signal end-of-body downstream + // without discarding anything the daemon still owes us. + errs <- conn.CloseWrite() + return + } + errs <- err + return + } + if _, err := conn.Write(msg.GetData()); err != nil { + errs <- err + return + } + } + }() + + if err := <-errs; err != nil { + return err + } + return <-errs +} + +// tcpPair returns the two ends of a real TCP connection. net.Pipe would be +// simpler but has no CloseWrite, and half-close is precisely what is under +// test. +func tcpPair(t *testing.T) (near, far *net.TCPConn, err error) { + t.Helper() + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, nil, err + } + defer ln.Close() + + type accepted struct { + conn net.Conn + err error + } + ch := make(chan accepted, 1) + go func() { + conn, err := ln.Accept() + ch <- accepted{conn: conn, err: err} + }() + + dialed, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + return nil, nil, err + } + + got := <-ch + if got.err != nil { + dialed.Close() + return nil, nil, got.err + } + + t.Cleanup(func() { + dialed.Close() + got.conn.Close() + }) + + return dialed.(*net.TCPConn), got.conn.(*net.TCPConn), nil +} + +// newStreamPair returns the two ends of an in-memory bidirectional stream with +// gRPC's semantics: a receive returns io.EOF once the peer has closed its send +// direction and every message already sent has been delivered, and message +// bytes are copied on send, as marshalling would copy them. +func newStreamPair(ctx context.Context) (*clientStream, *serverStream) { + var ( + toServer = make(chan []byte, 16) + toClient = make(chan []byte, 16) + ) + + c := &clientStream{halfStream: halfStream{ctx: ctx, send: toServer, recv: toClient}} + s := &serverStream{halfStream: halfStream{ctx: ctx, send: toClient, recv: toServer}} + + return c, s +} + +type halfStream struct { + ctx context.Context + send chan []byte + recv chan []byte + once sync.Once +} + +func (h *halfStream) SendMsg(m any) error { + msg, ok := m.(*ByteMessage) + if !ok { + return errUnexpectedMessage + } + + data := make([]byte, len(msg.GetData())) + copy(data, msg.GetData()) + + select { + case h.send <- data: + return nil + case <-h.ctx.Done(): + return h.ctx.Err() + } +} + +func (h *halfStream) RecvMsg(m any) error { + msg, ok := m.(*ByteMessage) + if !ok { + return errUnexpectedMessage + } + + select { + case data, open := <-h.recv: + if !open { + return io.EOF + } + msg.Data = data + return nil + case <-h.ctx.Done(): + return h.ctx.Err() + } +} + +func (h *halfStream) closeSend() { + h.once.Do(func() { close(h.send) }) +} + +func (h *halfStream) Context() context.Context { return h.ctx } + +type clientStream struct { + halfStream +} + +func (c *clientStream) Send(m *ByteMessage) error { return c.SendMsg(m) } + +func (c *clientStream) Recv() (*ByteMessage, error) { + msg := &ByteMessage{} + if err := c.RecvMsg(msg); err != nil { + return nil, err + } + return msg, nil +} + +// serverStream stands in for the generated Registry_ProxyServer. The embedded +// grpc.ServerStream covers the header and trailer methods the proxy never +// calls; the ones it does call are implemented above. +type serverStream struct { + halfStream + grpc.ServerStream +} + +func (s *serverStream) Send(m *ByteMessage) error { return s.halfStream.SendMsg(m) } + +func (s *serverStream) Recv() (*ByteMessage, error) { + msg := &ByteMessage{} + if err := s.halfStream.RecvMsg(msg); err != nil { + return nil, err + } + return msg, nil +} + +func (s *serverStream) SendMsg(m any) error { return s.halfStream.SendMsg(m) } + +func (s *serverStream) RecvMsg(m any) error { return s.halfStream.RecvMsg(m) } + +func (s *serverStream) Context() context.Context { return s.halfStream.Context() } + +var _ Registry_ProxyServer = (*serverStream)(nil) + +// blob is a stand-in for a layer: larger than the 32KiB copy buffers on both +// sides, so it crosses the stream as many messages. +func blob(n int) []byte { + b := make([]byte, n) + for i := range b { + b[i] = byte(i % 251) + } + return b +} + +func TestProxyServesAResponse(t *testing.T) { + body := blob(128 * 1024) + + reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Length", strconv.Itoa(len(body))) + _, _ = w.Write(body) + })) + t.Cleanup(reg.Close) + + client, _ := tunnel(t, reg) + + resp, err := client.Get("http://registry.invalid/v2/img/blobs/sha256:0") + if err != nil { + t.Fatalf("GET through the proxy: %v", err) + } + defer resp.Body.Close() + + got, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("reading the proxied body: %v", err) + } + + if len(got) != len(body) { + t.Fatalf("body length: got %d bytes, want %d", len(got), len(body)) + } + if string(got) != string(body) { + t.Error("body differs from what the registry served") + } +} + +func TestProxyDoesNotTruncateAResponseThatStallsMidBody(t *testing.T) { + body := blob(128 * 1024) + + reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Length", strconv.Itoa(len(body))) + half := len(body) / 2 + + _, _ = w.Write(body[:half]) + w.(http.Flusher).Flush() + + // A loaded runner -- the -race integration suite, say -- can leave a + // registry this long between chunks. Silence on the socket is not the + // end of a response. + time.Sleep(300 * time.Millisecond) + + _, _ = w.Write(body[half:]) + })) + t.Cleanup(reg.Close) + + client, _ := tunnel(t, reg) + + resp, err := client.Get("http://registry.invalid/v2/img/blobs/sha256:0") + if err != nil { + t.Fatalf("GET through the proxy: %v", err) + } + defer resp.Body.Close() + + got, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("reading a body the registry paused mid-way: %v (got %d of %d bytes)", err, len(got), len(body)) + } + if string(got) != string(body) { + t.Errorf("body differs from what the registry served: got %d bytes, want %d", len(got), len(body)) + } +} + +func TestProxyReusesOneConnectionForTwoRequests(t *testing.T) { + body := blob(4096) + + reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Length", strconv.Itoa(len(body))) + _, _ = w.Write(body) + })) + t.Cleanup(reg.Close) + + client, dials := tunnel(t, reg) + + for i := 0; i < 2; i++ { + if i > 0 { + // docker does its own work between a manifest fetch and the blob + // fetches that follow. A connection that only survives while + // requests arrive back-to-back is not a keep-alive connection. + time.Sleep(200 * time.Millisecond) + } + + resp, err := client.Get("http://registry.invalid/v2/img/manifests/latest") + if err != nil { + t.Fatalf("request %d through the proxy: %v", i+1, err) + } + + got, err := io.ReadAll(resp.Body) + resp.Body.Close() + if err != nil { + t.Fatalf("reading response %d: %v", i+1, err) + } + if string(got) != string(body) { + t.Fatalf("response %d differs from what the registry served", i+1) + } + } + + // docker holds keep-alive connections open across the many manifest and + // blob requests of one pull. Tearing the tunnel down after the first + // response makes every later request pay for a new stream, and strands any + // request already in flight on the old one. + if n := dials(); n != 1 { + t.Errorf("connections dialled: got %d, want 1 -- the tunnel did not survive the first response", n) + } +} + +func TestProxyPassesALargeRequestBodyThrough(t *testing.T) { + sent := blob(256 * 1024) + + var ( + mu sync.Mutex + received []byte + ) + + reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, err := io.ReadAll(r.Body) + mu.Lock() + received = b + mu.Unlock() + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + w.WriteHeader(http.StatusCreated) + })) + t.Cleanup(reg.Close) + + client, _ := tunnel(t, reg) + + resp, err := client.Post("http://registry.invalid/v2/img/blobs/uploads/", "application/octet-stream", bytes.NewReader(sent)) + if err != nil { + t.Fatalf("POST through the proxy: %v", err) + } + defer resp.Body.Close() + _, _ = io.Copy(io.Discard, resp.Body) + + if resp.StatusCode != http.StatusCreated { + t.Fatalf("status: got %d, want %d", resp.StatusCode, http.StatusCreated) + } + + mu.Lock() + defer mu.Unlock() + if len(received) != len(sent) { + t.Fatalf("request body length: got %d bytes, want %d", len(received), len(sent)) + } + if string(received) != string(sent) { + t.Error("request body differs from what was sent") + } +} diff --git a/api/services/registry/server.go b/api/services/registry/server.go index 0d7c8c118..7397f9015 100644 --- a/api/services/registry/server.go +++ b/api/services/registry/server.go @@ -1,17 +1,11 @@ package earthly_registry_v1 //nolint:revive import ( - "io" + "fmt" "net" "strings" - "time" - - "github.com/pkg/errors" - "golang.org/x/sync/errgroup" ) -const readDeadline = 50 * time.Millisecond - // NewServer creates and returns a new proxy server with a given host and client. func NewServer(addr string) *Server { return &Server{ @@ -25,126 +19,24 @@ type Server struct { UnimplementedRegistryServer } -type streamSource interface { - Send(*ByteMessage) error - Recv() (*ByteMessage, error) -} - -// NewStreamRW creates and returns a gRPC stream reader-writer that implements -// io.Reader & io.Writer as to utilize the gRPC stream with standard methods. -func NewStreamRW(stream streamSource) *StreamRW { - return &StreamRW{stream: stream} -} - -type StreamRW struct { - stream streamSource - last []byte -} - -// Write implements io.Writer. -func (s *StreamRW) Write(p []byte) (int, error) { - err := s.stream.Send(&ByteMessage{ - Data: p, - }) - if err != nil { - return 0, errors.Wrap(err, "failed to write data to client") - } - return len(p), nil -} - -// Read implements io.Reader. -func (s *StreamRW) Read(p []byte) (int, error) { - l := 0 - if len(s.last) > 0 { - l = copy(p, s.last) - } - - msg, err := s.stream.Recv() - if err != nil { - return 0, err - } - - s.last = msg.GetData() - n := copy(p, s.last) - s.last = s.last[n:] - - return n + l, nil -} - // Proxy requests sent via gRPC data stream to the embedded Docker registry and // pipe them back out through the stream again. This allows us to send HTTP // requests to the embedded registry without having to connect via some other // exposed server or port. +// +// One stream carries one connection from the client's local listener, for as +// long as the client keeps it: docker reuses a connection across the manifest +// and blob requests of a single pull, and each of those requests is answered +// on the stream that carried it. func (s *Server) Proxy(stream Registry_ProxyServer) error { - rw := NewStreamRW(stream) - addr := strings.ReplaceAll(s.addr, "0.0.0.0", "127.0.0.1") conn, err := net.Dial("tcp", addr) if err != nil { - return err - } - defer conn.Close() - - ctx := stream.Context() - eg, _ := errgroup.WithContext(ctx) - - eg.Go(func() error { - _, err = io.Copy(conn, rw) - if err != nil { - return errors.Wrap(err, "failed to copy from stream to host") - } - return nil - }) - - eg.Go(func() error { - _, err = CopyWithDeadline(conn, rw) - if err != nil { - return errors.Wrap(err, "failed to copy from host to stream") - } - return nil - }) - - err = eg.Wait() - if err != nil { - return errors.Wrap(err, "failed to wait") - } - - return nil -} - -// CopyWithDeadline copies data from a net.Conn using a read deadline. The -// process will fail with a timeout error if no data is read for the defined -// period. -func CopyWithDeadline(conn net.Conn, w io.Writer) (int64, error) { - var ( - t = int64(0) - buf = make([]byte, 32*1024) - ) - for { - err := conn.SetReadDeadline(time.Now().Add(readDeadline)) - if err != nil { - return t, err - } - n, err := conn.Read(buf) - if err != nil { - if errors.Is(err, io.EOF) || isNetTimeout(err) { - break - } - return t, err - } - n, err = w.Write(buf[0:n]) - t += int64(n) - if err != nil { - return t, err - } + return fmt.Errorf("dial embedded registry at %s: %w", addr, err) } - return t, nil -} -func isNetTimeout(err error) bool { - if netErr, ok := err.(net.Error); ok && netErr.Timeout() { - return true - } - return false + // The stream is closed by returning from this handler, so there is no send + // direction for the copy to close on its own. + return Copy(stream.Context(), conn, stream, nil) }