From 6b8278040d5ff8447281744d3eeb0b92ed95d156 Mon Sep 17 00:00:00 2001 From: Copybara Service Date: Thu, 22 Oct 2026 07:17:48 -0700 Subject: [PATCH] Stop failing compressed polls and stalling on empty frames. PiperOrigin-RevId: 1004753414 --- fleetspeak/src/server/https/compression.go | 4 + .../src/server/https/compression_test.go | 124 ++++++ fleetspeak/src/server/https/https.go | 14 +- fleetspeak/src/server/https/https_test.go | 366 ++++++++++++++++++ fleetspeak/src/server/https/message_server.go | 7 + 5 files changed, 514 insertions(+), 1 deletion(-) create mode 100644 fleetspeak/src/server/https/compression_test.go diff --git a/fleetspeak/src/server/https/compression.go b/fleetspeak/src/server/https/compression.go index b39d8de3..11b8beca 100644 --- a/fleetspeak/src/server/https/compression.go +++ b/fleetspeak/src/server/https/compression.go @@ -32,6 +32,10 @@ func (h *CompressionHandler) ServeHTTP(res http.ResponseWriter, req *http.Reques defer zr.Close() req.Body = zr req.Header.Del("Content-Encoding") + // The length of the request describes the compressed body. The length of + // the decompressed body is unknown. + req.ContentLength = -1 + req.Header.Del("Content-Length") default: http.Error(res, fmt.Sprintf("unsupported content encoding: %s", encoding), http.StatusUnsupportedMediaType) return diff --git a/fleetspeak/src/server/https/compression_test.go b/fleetspeak/src/server/https/compression_test.go new file mode 100644 index 00000000..3b8e6bfd --- /dev/null +++ b/fleetspeak/src/server/https/compression_test.go @@ -0,0 +1,124 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package https + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + + fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak" +) + +func TestCompressionHandler(t *testing.T) { + data := []byte(strings.Repeat("fleetspeak", 100)) + compressed := compress(t, data, fspb.CompressionAlgorithm_COMPRESSION_DEFLATE) + + testCases := []struct { + name string + encoding string + body []byte + // wantStatus is the status of the response. It is http.StatusOK if the + // request was passed to the wrapped handler. + wantStatus int + wantBody []byte + wantContentLength int64 + }{ + { + name: "uncompressed", + body: data, + wantStatus: http.StatusOK, + wantBody: data, + wantContentLength: int64(len(data)), + }, + { + // The length of the decompressed body is unknown. + name: "deflate", + encoding: "deflate", + body: compressed, + wantStatus: http.StatusOK, + wantBody: data, + wantContentLength: -1, + }, + { + name: "invalid deflate data", + encoding: "deflate", + body: data, + wantStatus: http.StatusBadRequest, + }, + { + name: "unsupported encoding", + encoding: "gzip", + body: data, + wantStatus: http.StatusUnsupportedMediaType, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var gotReq *http.Request + var gotBody []byte + h := &CompressionHandler{ + Wrapped: http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) { + gotReq = req + var err error + if gotBody, err = io.ReadAll(req.Body); err != nil { + t.Errorf("Error reading request body: %v", err) + } + }), + } + req := httptest.NewRequest(http.MethodPost, "/message", bytes.NewReader(tc.body)) + req.Header.Set("Content-Length", strconv.Itoa(len(tc.body))) + if tc.encoding != "" { + req.Header.Set("Content-Encoding", tc.encoding) + } + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != tc.wantStatus { + t.Errorf("ServeHTTP() returned status %d, want %d", rec.Code, tc.wantStatus) + } + if gotReq == nil { + if tc.wantStatus == http.StatusOK { + t.Error("ServeHTTP() didn't call the wrapped handler") + } + return + } + if tc.wantStatus != http.StatusOK { + t.Fatal("ServeHTTP() called the wrapped handler") + } + if !bytes.Equal(gotBody, tc.wantBody) { + t.Errorf("Wrapped handler got body %q, want %q", gotBody, tc.wantBody) + } + if gotReq.ContentLength != tc.wantContentLength { + t.Errorf("Wrapped handler got ContentLength %d, want %d", gotReq.ContentLength, tc.wantContentLength) + } + // The headers have to describe the body that the wrapped handler gets. + if got := gotReq.Header.Get("Content-Encoding"); got != "" { + t.Errorf("Wrapped handler got Content-Encoding header %q, want none", got) + } + wantLengthHeader := "" + if tc.wantContentLength >= 0 { + wantLengthHeader = strconv.FormatInt(tc.wantContentLength, 10) + } + if got := gotReq.Header.Get("Content-Length"); got != wantLengthHeader { + t.Errorf("Wrapped handler got Content-Length header %q, want %q", got, wantLengthHeader) + } + }) + } +} diff --git a/fleetspeak/src/server/https/https.go b/fleetspeak/src/server/https/https.go index df937b7e..4107d908 100644 --- a/fleetspeak/src/server/https/https.go +++ b/fleetspeak/src/server/https/https.go @@ -250,7 +250,11 @@ var bufferPool = sync.Pool{ }, } +// getBuffer returns an empty buffer with a capacity of at least size bytes. +// Buffers of up to maxSize bytes are taken from bufferPool. A negative size is +// treated as 0. func getBuffer(size, maxSize int) *bytes.Buffer { + size = max(size, 0) if size > maxSize { return bytes.NewBuffer(make([]byte, 0, size)) } @@ -266,12 +270,20 @@ func putBuffer(b *bytes.Buffer, maxSize int) { } } +// readWrappedContactData reads a WrappedContactData from r and returns it, +// together with the number of bytes read. +// +// A non-negative size is the exact size of the WrappedContactData, e.g. the +// length of a frame on a streaming connection, and exactly that many bytes are +// read from r. A size of 0 denotes an empty WrappedContactData. A negative size +// means that the size is unknown, e.g. for a decompressed request body, and r +// is read until EOF. func readWrappedContactData(r io.Reader, size, maxReusableBufferSize int) (*fspb.WrappedContactData, int, error) { b := getBuffer(size, maxReusableBufferSize) defer putBuffer(b, maxReusableBufferSize) var err error - if size > 0 { + if size >= 0 { _, err = io.CopyN(b, r, int64(size)) } else { _, err = b.ReadFrom(r) diff --git a/fleetspeak/src/server/https/https_test.go b/fleetspeak/src/server/https/https_test.go index b83b6588..e031158f 100644 --- a/fleetspeak/src/server/https/https_test.go +++ b/fleetspeak/src/server/https/https_test.go @@ -35,6 +35,7 @@ import ( "net/http" "net/url" "path" + "slices" "strings" "sync" "testing" @@ -477,6 +478,371 @@ func TestStreaming(t *testing.T) { } } +// marshalWrapped returns a serialized WrappedContactData with the given client +// labels, wrapping a ContactData that only has its client clock set. +func marshalWrapped(t *testing.T, labels ...string) []byte { + t.Helper() + cd, err := proto.Marshal(&fspb.ContactData{ClientClock: db.NowProto()}) + if err != nil { + t.Fatal(err) + } + b, err := proto.Marshal(&fspb.WrappedContactData{ContactData: cd, ClientLabels: labels}) + if err != nil { + t.Fatal(err) + } + return b +} + +// compress returns b compressed with the given algorithm, as the client sends +// it. +func compress(t *testing.T, b []byte, compression fspb.CompressionAlgorithm) []byte { + t.Helper() + var buf bytes.Buffer + bw := https.CompressingWriter(&buf, compression) + if _, err := bw.Write(b); err != nil { + t.Fatal(err) + } + if err := bw.Close(); err != nil { + t.Fatal(err) + } + return buf.Bytes() +} + +// postMessage sends body to the /message endpoint of the server at addr, +// announcing the given compression. If hideLength is set, the body is sent +// without a Content-Length header, i.e. chunked. +func postMessage(t *testing.T, cl *http.Client, addr string, body []byte, compression fspb.CompressionAlgorithm, hideLength bool) *http.Response { + t.Helper() + u := url.URL{Scheme: "https", Host: addr, Path: "/message"} + var r io.Reader = bytes.NewReader(body) + if hideLength { + // Hide the type of the reader, so that its length is unknown. + r = io.MultiReader(r) + } + req, err := http.NewRequest(http.MethodPost, u.String(), r) + if err != nil { + t.Fatal(err) + } + if hideLength { + req.ContentLength = -1 + } + req.Close = true + https.SetContentEncoding(req.Header, compression) + resp, err := cl.Do(req) + if err != nil { + t.Fatalf("POST to /message failed: %v", err) + } + t.Cleanup(func() { resp.Body.Close() }) + return resp +} + +// checkPollResponse checks that resp is a successful response to a poll. +func checkPollResponse(t *testing.T, resp *http.Response) { + t.Helper() + b, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("Error reading response body: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("Poll returned status %d (%q), want %d", resp.StatusCode, b, http.StatusOK) + } + var cd fspb.ContactData + if err := proto.Unmarshal(b, &cd); err != nil { + t.Fatalf("Unable to parse returned data as ContactData: %v", err) + } + if cd.SequencingNonce == 0 { + t.Error("Expected SequencingNonce in returned ContactData") + } +} + +// checkClientLabels checks that the client with the given id exists in the +// datastore and has exactly the given labels. +func checkClientLabels(t *testing.T, ds *sqlite.Datastore, id common.ClientID, want []string) { + t.Helper() + cd, err := ds.GetClientData(context.Background(), id) + if err != nil { + t.Fatalf("Error getting client data after poll: %v", err) + } + var got []string + for _, l := range cd.Labels { + got = append(got, l.GetLabel()) + } + slices.Sort(got) + want = slices.Sorted(slices.Values(want)) + if !slices.Equal(got, want) { + t.Errorf("Client has labels %q, want %q", got, want) + } +} + +func TestReadWrappedContactData(t *testing.T) { + wcdBytes := marshalWrapped(t, "linux", "test") + wcd := &fspb.WrappedContactData{} + if err := proto.Unmarshal(wcdBytes, wcd); err != nil { + t.Fatal(err) + } + const next = "next frame" + + testCases := []struct { + name string + input []byte + size int + want *fspb.WrappedContactData + wantReadBytes int + wantUnread string + wantErr bool + }{ + { + name: "exact size", + input: append(bytes.Clone(wcdBytes), next...), + size: len(wcdBytes), + want: wcd, + wantReadBytes: len(wcdBytes), + wantUnread: next, + }, + { + name: "zero size", + input: []byte(next), + size: 0, + want: &fspb.WrappedContactData{}, + wantReadBytes: 0, + wantUnread: next, + }, + { + name: "unknown size", + input: wcdBytes, + size: -1, + want: wcd, + wantReadBytes: len(wcdBytes), + }, + { + name: "empty input of unknown size", + input: nil, + size: -1, + want: &fspb.WrappedContactData{}, + wantReadBytes: 0, + }, + { + name: "truncated input", + input: wcdBytes[:len(wcdBytes)-1], + size: len(wcdBytes), + wantReadBytes: len(wcdBytes) - 1, + wantErr: true, + }, + { + name: "invalid input", + input: []byte(next), + size: len(next), + wantReadBytes: len(next), + wantErr: true, + }, + } + for _, tc := range testCases { + // Cover both pooled and freshly allocated buffers. + for _, maxReusableBufferSize := range []int{0, 1024 * 1024} { + t.Run(fmt.Sprintf("%s/maxReusableBufferSize=%d", tc.name, maxReusableBufferSize), func(t *testing.T) { + r := bytes.NewReader(tc.input) + got, readBytes, err := readWrappedContactData(r, tc.size, maxReusableBufferSize) + if (err != nil) != tc.wantErr { + t.Errorf("readWrappedContactData(size=%d) returned error %v, want error: %t", tc.size, err, tc.wantErr) + } + if !tc.wantErr && !proto.Equal(got, tc.want) { + t.Errorf("readWrappedContactData(size=%d) = %v, want %v", tc.size, got, tc.want) + } + if readBytes != tc.wantReadBytes { + t.Errorf("readWrappedContactData(size=%d) read %d bytes, want %d", tc.size, readBytes, tc.wantReadBytes) + } + unread, err := io.ReadAll(r) + if err != nil { + t.Fatal(err) + } + if string(unread) != tc.wantUnread { + t.Errorf("readWrappedContactData(size=%d) left %q unread, want %q", tc.size, unread, tc.wantUnread) + } + }) + } + } +} + +func TestCompressedPoll(t *testing.T) { + s, ds, addr := makeServer(t, "CompressedPoll", nil, nil) + defer s.Stop() + + testCases := []struct { + name string + labels []string + hideLength bool + }{ + { + // The compressed body is much shorter than the WrappedContactData. + name: "compressible", + labels: []string{"linux", strings.Repeat("compressible", 100)}, + }, + { + // With the zlib header and checksum, the compressed body is longer than + // the WrappedContactData. + name: "small", + labels: []string{"linux"}, + }, + { + name: "without content length", + labels: []string{"linux", strings.Repeat("compressible", 100)}, + hideLength: true, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + id, cl, _, _ := makeClient(t, true) + body := compress(t, marshalWrapped(t, tc.labels...), fspb.CompressionAlgorithm_COMPRESSION_DEFLATE) + resp := postMessage(t, cl, addr, body, fspb.CompressionAlgorithm_COMPRESSION_DEFLATE, tc.hideLength) + checkPollResponse(t, resp) + // The labels show that the server read the whole WrappedContactData. + checkClientLabels(t, ds, id, tc.labels) + }) + } +} + +func TestPollWithoutContentLength(t *testing.T) { + s, ds, addr := makeServer(t, "PollWithoutContentLength", nil, nil) + defer s.Stop() + id, cl, _, _ := makeClient(t, true) + + labels := []string{"linux", "test"} + resp := postMessage(t, cl, addr, marshalWrapped(t, labels...), fspb.CompressionAlgorithm_COMPRESSION_NONE, true) + checkPollResponse(t, resp) + checkClientLabels(t, ds, id, labels) +} + +func TestPollTooLarge(t *testing.T) { + ctx := context.Background() + s, ds, addr := makeServer(t, "PollTooLarge", nil, nil) + defer s.Stop() + + // A valid WrappedContactData of MaxContactSize+1 bytes: the contact_data + // field takes 1 byte for its tag and 4 bytes for its length. + tooLarge, err := proto.Marshal(&fspb.WrappedContactData{ContactData: make([]byte, MaxContactSize+1-5)}) + if err != nil { + t.Fatal(err) + } + if len(tooLarge) != MaxContactSize+1 { + t.Fatalf("WrappedContactData has %d bytes, want %d", len(tooLarge), MaxContactSize+1) + } + + // The sizes of these bodies are unknown until the server has read them. + testCases := []struct { + name string + compression fspb.CompressionAlgorithm + }{ + { + name: "uncompressed", + compression: fspb.CompressionAlgorithm_COMPRESSION_NONE, + }, + { + name: "deflate", + compression: fspb.CompressionAlgorithm_COMPRESSION_DEFLATE, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + id, cl, _, _ := makeClient(t, true) + body := compress(t, tooLarge, tc.compression) + resp := postMessage(t, cl, addr, body, tc.compression, true) + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("Poll returned status %d, want %d", resp.StatusCode, http.StatusBadRequest) + } + if _, err := ds.GetClientData(ctx, id); err == nil { + t.Error("Client exists in the datastore after a poll that was too large") + } + }) + } +} + +func TestStreamingEmptyContacts(t *testing.T) { + s, _, addr := makeServer(t, "StreamingEmptyContacts", nil, nil) + defer s.Stop() + + testCases := []struct { + name string + compression fspb.CompressionAlgorithm + }{ + { + name: "uncompressed", + compression: fspb.CompressionAlgorithm_COMPRESSION_NONE, + }, + { + name: "deflate", + compression: fspb.CompressionAlgorithm_COMPRESSION_DEFLATE, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + // The stream stays open, so the server must not wait for its end. + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _, cl, _, _ := makeClient(t, true) + + pr, pw := io.Pipe() + defer pr.Close() + // Ending the request body is the only way to unblock the client if the + // server doesn't respond. + stop := context.AfterFunc(ctx, func() { pr.CloseWithError(context.Cause(ctx)) }) + defer stop() + bw := https.CompressingWriter(pw, tc.compression) + defer bw.Close() + emptyFrame := binary.AppendUvarint(nil, 0) + + initialContactWritten := make(chan struct{}) + go func() { + defer close(initialContactWritten) + // An empty WrappedContactData is a valid, though atypical, initial + // contact. It has to be written during the http.Client.Do call below, + // because the server writes headers at the end of the first exchange. + binary.Write(bw, binary.LittleEndian, magic) + if _, err := bw.Write(emptyFrame); err != nil { + t.Error(err) + } + bw.Flush() + }() + + u := url.URL{Scheme: "https", Host: addr, Path: "/streaming-message"} + req, err := http.NewRequestWithContext(ctx, http.MethodPost, u.String(), pr) + if err != nil { + t.Fatal(err) + } + req.ContentLength = -1 + req.Close = true + req.Header.Set("Expect", "100-continue") + https.SetContentEncoding(req.Header, tc.compression) + resp, err := cl.Do(req) + if err != nil { + t.Fatalf("Streaming post failed: %v", err) + } + defer resp.Body.Close() + body := bufio.NewReader(resp.Body) + cd, err := readContact(body) + if err != nil { + t.Fatalf("Error reading ContactData of the initial exchange: %v", err) + } + if cd.AckIndex != 0 { + t.Errorf("AckIndex of initial exchange should be unset, got %d", cd.AckIndex) + } + <-initialContactWritten + + // Empty WrappedContactData on the stream are acked like any other. + if _, err := bw.Write(emptyFrame); err != nil { + t.Fatal(err) + } + bw.Flush() + cd, err = readContact(body) + if err != nil { + t.Fatalf("Error reading ack of an empty WrappedContactData: %v", err) + } + if cd.AckIndex != 1 { + t.Errorf("Received ack for contact %d, but expected 1", cd.AckIndex) + } + }) + } +} + func TestHeaderNormalPoll(t *testing.T) { ctx := context.Background() clientCertHeader := "ssl-client-cert" diff --git a/fleetspeak/src/server/https/message_server.go b/fleetspeak/src/server/https/message_server.go index 9758af0d..e71ce7b8 100644 --- a/fleetspeak/src/server/https/message_server.go +++ b/fleetspeak/src/server/https/message_server.go @@ -154,6 +154,13 @@ func (s messageServer) ServeHTTP(res http.ResponseWriter, req *http.Request) { wcd, readBytes, err := readWrappedContactData(req.Body, int(req.ContentLength), s.p.MaxReusableBufferSize) pi.ReadTime = time.Since(st) pi.ReadBytes = readBytes + // The length of compressed and chunked bodies is only known once they have + // been read. + if readBytes > MaxContactSize { + pi.Status = http.StatusBadRequest + http.Error(res, fmt.Sprintf("body can't be larger than %v bytes", MaxContactSize), pi.Status) + return + } if err != nil { pi.Status = http.StatusBadRequest http.Error(res, err.Error(), pi.Status)