From 9e97f98d12408cfc8fb355bbf610c5e06fac8eb8 Mon Sep 17 00:00:00 2001 From: AdrianArnautu Date: Thu, 9 Mar 2023 10:38:39 +0200 Subject: [PATCH] This change allows the HTTP clients to consume and send gzip compressed response and request body. (#5696) It is available for the following REST API endpoints: - GET & POST HTTP methods on /v0/data & /v1/data endpoints - POST HTTP method on /v1/compile endpoint HTTP clients can optionally: - send 'Accept-Encoding: gzip' header and expect a gzip compressed body and a Content-Encoding: gzip response header. The server will send the content encoded as gzip only after a threshold defined by server.encoding.gzip.min_length (default value is 1024). If the size is below the threshold, the body is not compressed - send 'Content-Encoding: gzip' header and a gzip compressed body and expect the server to correctly interpret the request Fixes #5310 Signed-off-by: aarnautu --- config/config.go | 5 +- config/config_test.go | 16 + docs/content/configuration.md | 16 + docs/content/rest-api.md | 13 + plugins/server/encoding/config.go | 91 ++++++ plugins/server/encoding/config_test.go | 105 +++++++ runtime/logging.go | 58 +++- runtime/logging_test.go | 165 +++++++++- server/handlers/compress.go | 194 ++++++++++++ server/handlers/compress_test.go | 182 +++++++++++ server/server.go | 70 ++++- server/server_test.go | 402 +++++++++++++++++++++++++ 12 files changed, 1296 insertions(+), 21 deletions(-) create mode 100644 plugins/server/encoding/config.go create mode 100644 plugins/server/encoding/config_test.go create mode 100644 server/handlers/compress.go create mode 100644 server/handlers/compress_test.go diff --git a/config/config.go b/config/config.go index 57c9a5b641..51931e440f 100644 --- a/config/config.go +++ b/config/config.go @@ -35,7 +35,10 @@ type Config struct { NDBuiltinCache bool `json:"nd_builtin_cache,omitempty"` PersistenceDirectory *string `json:"persistence_directory,omitempty"` DistributedTracing json.RawMessage `json:"distributed_tracing,omitempty"` - Storage *struct { + Server *struct { + Encoding json.RawMessage `json:"encoding,omitempty"` + } `json:"server,omitempty"` + Storage *struct { Disk json.RawMessage `json:"disk,omitempty"` } `json:"storage,omitempty"` } diff --git a/config/config_test.go b/config/config_test.go index f743cdba4d..43a508ff10 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -196,6 +196,14 @@ func TestActiveConfig(t *testing.T) { "plugins": { "some-plugin": {} }, + "server": { + "encoding": { + "gzip": { + "min_length": 1024, + "compression_level": 1 + } + } + }, "discovery": {"name": "config"}` serviceObj := `"services": { @@ -249,6 +257,14 @@ func TestActiveConfig(t *testing.T) { "plugins": { "some-plugin": {} }, + "server": { + "encoding": { + "gzip": { + "min_length": 1024, + "compression_level": 1 + } + } + }, "default_authorization_decision": "/system/authz/allow", "default_decision": "/system/main", "discovery": {"name": "config"}`, version.Version) diff --git a/docs/content/configuration.md b/docs/content/configuration.md index bcd9eb45f1..09d7aaa93e 100644 --- a/docs/content/configuration.md +++ b/docs/content/configuration.md @@ -76,6 +76,12 @@ distributed_tracing: service_name: opa sample_percentage: 50 encryption: "off" + +server: + encoding: + gzip: + min_length: 1024, + compression_level: 9 ``` #### Environment Variable Substitution @@ -925,3 +931,13 @@ with data put into the configured `directory`. | `storage.disk.badger` | `string` | No (default: empty) | "Superflags" passed to Badger allowing to modify advanced options. | See [the docs on disk storage](../misc-disk/) for details about the settings. + +### Server + +The `server` configuration sets the gzip compression settings for `/v0/data`, `/v1/data` and `/v1/compile` HTTP `POST` endpoints +The gzip compression settings are used when the client sends `Accept-Encoding: gzip` + +| Field | Type | Required | Description | +|------------------------------------------|-------|---------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `server.encoding.gzip.min_length` | `int` | No, (default: 1024) | Specifies the minimum length of the response to compress | +| `server.encoding.gzip.compression_level` | `int` | No, (default: 9) | Specifies the compression level. Accepted values: a value of either 0 (no compression), 1 (best speed, lowest compression) or 9 (slowest, best compression). See https://pkg.go.dev/compress/flate#pkg-constants | diff --git a/docs/content/rest-api.md b/docs/content/rest-api.md index f8f9734008..69c41f2b41 100644 --- a/docs/content/rest-api.md +++ b/docs/content/rest-api.md @@ -732,6 +732,10 @@ The path separator is used to access values inside object and array documents. I - **instrument** - Instrument query evaluation and return a superset of performance metrics in addition to result. See [Performance Metrics](#performance-metrics) for more detail. - **strict-builtin-errors** - Treat built-in function call errors as fatal and return an error immediately. +#### Request Headers + +- **Accept-Encoding: gzip**: Indicates the server should respond with a gzip encoded body. The server will send the compressed response only if its length is above `server.encoding.gzip.min_length` value. See the configuration section + #### Status Codes - **200** - no error @@ -820,6 +824,8 @@ The request body contains an object that specifies a value for [The input Docume #### Request Headers - **Content-Type: application/x-yaml**: Indicates the request body is a YAML encoded object. +- **Content-Encoding: gzip**: Indicates the request body is a gzip encoded object. +- **Accept-Encoding: gzip**: Indicates the server should respond with a gzip encoded body. The server will send the compressed response only if its length is above `server.encoding.gzip.min_length` value. See the configuration section #### Query Parameters @@ -939,6 +945,8 @@ array documents. #### Request Headers - **Content-Type: application/x-yaml**: Indicates the request body is a YAML encoded object. +- **Content-Encoding: gzip**: Indicates the request body is a gzip encoded object. +- **Accept-Encoding: gzip**: Indicates the server should respond with a gzip encoded body. The server will send the compressed response only if its length is above `server.encoding.gzip.min_length` value. See the configuration section #### Query Parameters @@ -1290,6 +1298,11 @@ Compile API requests contain the following fields: | `options` | `object[string, any]` | No | Additional options to use during partial evaluation. Only `disableInlining` option is supported. (default: undefined). | | `unknowns` | `array[string]` | No | The terms to treat as unknown during partial evaluation (default: `["input"]`]). | +### Request Headers + +- **Content-Encoding: gzip**: Indicates the request body is a gzip encoded object. +- **Accept-Encoding: gzip**: Indicates the server should respond with a gzip encoded body. The server will send the compressed response only if its length is above `server.encoding.gzip.min_length` value + #### Query Parameters - **pretty** - If parameter is `true`, response will formatted for humans. diff --git a/plugins/server/encoding/config.go b/plugins/server/encoding/config.go new file mode 100644 index 0000000000..4da12d0c07 --- /dev/null +++ b/plugins/server/encoding/config.go @@ -0,0 +1,91 @@ +package encoding + +import ( + "compress/gzip" + "fmt" + + "github.com/open-policy-agent/opa/util" +) + +var defaultGzipMinLength = 1024 +var defaultGzipCompressionLevel = gzip.BestCompression + +// Config represents the configuration for the Server.Encoding settings +type Config struct { + Gzip *Gzip `json:"gzip,omitempty"` +} + +// Gzip represents the configuration for the Server.Encoding.Gzip settings +type Gzip struct { + MinLength *int `json:"min_length,omitempty"` // the minimum length of a response that will be gzipped + CompressionLevel *int `json:"compression_level,omitempty"` // the compression level for gzip +} + +// ConfigBuilder assists in the construction of the plugin configuration. +type ConfigBuilder struct { + raw []byte +} + +// NewConfigBuilder returns a new ConfigBuilder to build and parse the server config +func NewConfigBuilder() *ConfigBuilder { + return &ConfigBuilder{} +} + +// WithBytes sets the raw server config +func (b *ConfigBuilder) WithBytes(config []byte) *ConfigBuilder { + b.raw = config + return b +} + +// Parse returns a valid Config object with defaults injected. +func (b *ConfigBuilder) Parse() (*Config, error) { + if b.raw == nil { + defaultConfig := &Config{ + Gzip: &Gzip{ + MinLength: &defaultGzipMinLength, + CompressionLevel: &defaultGzipCompressionLevel, + }, + } + return defaultConfig, nil + } + + var result Config + + if err := util.Unmarshal(b.raw, &result); err != nil { + return nil, err + } + + return &result, result.validateAndInjectDefaults() +} + +func (c *Config) validateAndInjectDefaults() error { + if c.Gzip == nil { + c.Gzip = &Gzip{ + MinLength: &defaultGzipMinLength, + CompressionLevel: &defaultGzipCompressionLevel, + } + } + if c.Gzip.MinLength == nil { + c.Gzip.MinLength = &defaultGzipMinLength + } + + if c.Gzip.CompressionLevel == nil { + c.Gzip.CompressionLevel = &defaultGzipCompressionLevel + } + + if *c.Gzip.MinLength <= 0 { + return fmt.Errorf("invalid value for server.encoding.gzip.min_length field, should be a positive number") + } + + acceptedCompressionLevels := map[int]bool{ + gzip.NoCompression: true, + gzip.BestSpeed: true, + gzip.BestCompression: true, + } + _, compressionLevelAccepted := acceptedCompressionLevels[*c.Gzip.CompressionLevel] + if !compressionLevelAccepted { + return fmt.Errorf("invalid value for server.encoding.gzip.compression_level field, accepted values are 0, 1 or 9") + } + + return nil +} diff --git a/plugins/server/encoding/config_test.go b/plugins/server/encoding/config_test.go new file mode 100644 index 0000000000..d8aa7482a3 --- /dev/null +++ b/plugins/server/encoding/config_test.go @@ -0,0 +1,105 @@ +package encoding + +import ( + "fmt" + "testing" +) + +func TestConfigValidation(t *testing.T) { + tests := []struct { + input string + wantErr bool + }{ + { + input: `{}`, + wantErr: false, + }, + { + input: `{"gzip": {"min_length": "not-a-number"}}`, + wantErr: true, + }, + { + input: `{"gzip": {min_length": 42}}`, + wantErr: false, + }, + { + input: `{"gzip":{"min_length": "42"}}`, + wantErr: true, + }, + { + input: `{"gzip":{"min_length": 0}}`, + wantErr: true, + }, + { + input: `{"gzip":{"min_length": -10}}`, + wantErr: true, + }, + { + input: `{"gzip":{"random_key": 0}}`, + wantErr: false, + }, + { + input: `{"gzip": {"min_length": -10, "compression_level": 13}}`, + wantErr: true, + }, + { + input: `{"gzip":{"compression_level": "not-an-number"}}`, + wantErr: true, + }, + { + input: `{"gzip":{"compression_level": 1}}`, + wantErr: false, + }, + { + input: `{"gzip":{"compression_level": 13}}`, + wantErr: true, + }, + { + input: `{"gzip":{"min_length": 42, "compression_level": 9}}`, + wantErr: false, + }, + } + + for i, test := range tests { + t.Run(fmt.Sprintf("TestConfigValidation_case_%d", i), func(t *testing.T) { + _, err := NewConfigBuilder().WithBytes([]byte(test.input)).Parse() + if err != nil && !test.wantErr { + t.Fail() + } + if err == nil && test.wantErr { + t.Fail() + } + }) + } +} + +func TestConfigValue(t *testing.T) { + tests := []struct { + input string + minLengthExpectedValue int + compressionLevelExpectedValue int + }{ + { + input: `{}`, + minLengthExpectedValue: 1024, + compressionLevelExpectedValue: 9, + }, + { + input: `{"gzip":{"min_length": 42, "compression_level": 1}}`, + minLengthExpectedValue: 42, + compressionLevelExpectedValue: 1, + }, + } + + for i, test := range tests { + t.Run(fmt.Sprintf("TestConfigValue_case_%d", i), func(t *testing.T) { + config, err := NewConfigBuilder().WithBytes([]byte(test.input)).Parse() + if err != nil { + t.Fail() + } + if *config.Gzip.MinLength != test.minLengthExpectedValue || *config.Gzip.CompressionLevel != test.compressionLevelExpectedValue { + t.Fail() + } + }) + } +} diff --git a/runtime/logging.go b/runtime/logging.go index f224dc0246..f18450ca5b 100644 --- a/runtime/logging.go +++ b/runtime/logging.go @@ -6,6 +6,7 @@ package runtime import ( "bytes" + "compress/gzip" "io" "net/http" "strings" @@ -73,8 +74,26 @@ func (h *LoggingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { bs, r.Body, err = readBody(r.Body) } if err == nil { - fields["req_body"] = string(bs) - } else { + if gzipReceived(r.Header) { + // the request is compressed + var gzReader *gzip.Reader + var plainOutput []byte + reader := bytes.NewReader(bs) + gzReader, err = gzip.NewReader(reader) + if err == nil { + plainOutput, err = io.ReadAll(gzReader) + if err == nil { + defer gzReader.Close() + fields["req_body"] = string(plainOutput) + } + } + } else { + fields["req_body"] = string(bs) + } + } + + // err can be thrown on different statements + if err != nil { fields["err"] = err } @@ -127,6 +146,21 @@ func (h *LoggingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { // metrics endpoint does so when the client accepts it (e.g. prometheus) fields["resp_body"] = "[compressed payload]" + case gzipAccepted(r.Header) && gzipReceived(w.Header()) && (isDataEndpoint(r) || isCompileEndpoint(r)): + // data and compile endpoints might compress the response + gzReader, gzErr := gzip.NewReader(recorder.buf) + if gzErr == nil { + plainOutput, readErr := io.ReadAll(gzReader) + if readErr == nil { + defer gzReader.Close() + fields["resp_body"] = string(plainOutput) + } else { + h.logger.Error("Failed to decompressed the payload: %v", readErr.Error()) + } + } else { + h.logger.Error("Failed to read the compressed payload: %v", gzErr.Error()) + } + default: fields["resp_body"] = recorder.buf.String() } @@ -148,6 +182,18 @@ func gzipAccepted(header http.Header) bool { return false } +func gzipReceived(header http.Header) bool { + a := header.Get("Content-Encoding") + parts := strings.Split(a, ",") + for _, part := range parts { + part = strings.TrimSpace(part) + if part == "gzip" || strings.HasPrefix(part, "gzip;") { + return true + } + } + return false +} + func isPprofEndpoint(req *http.Request) bool { return strings.HasPrefix(req.URL.Path, "/debug/pprof/") } @@ -156,6 +202,14 @@ func isMetricsEndpoint(req *http.Request) bool { return strings.HasPrefix(req.URL.Path, "/metrics") } +func isDataEndpoint(req *http.Request) bool { + return strings.HasPrefix(req.URL.Path, "/v1/data") || strings.HasPrefix(req.URL.Path, "/v0/data") +} + +func isCompileEndpoint(req *http.Request) bool { + return strings.HasPrefix(req.URL.Path, "/v1/compile") +} + type recorder struct { logger logging.Logger inner http.ResponseWriter diff --git a/runtime/logging_test.go b/runtime/logging_test.go index 5a2c784de2..6a823ac388 100644 --- a/runtime/logging_test.go +++ b/runtime/logging_test.go @@ -6,7 +6,11 @@ package runtime import ( + "bytes" + "compress/gzip" "context" + "io" + "log" "net/http" "net/http/httptest" "net/url" @@ -41,6 +45,30 @@ func TestValidateGzipHeader(t *testing.T) { } } +func TestValidateReceivedGzipHeader(t *testing.T) { + + httpHeader := http.Header{} + httpHeader.Set("Content-Encoding", "*/*") + if result, expected := gzipReceived(httpHeader), false; result != expected { + t.Errorf("Expected %v but got: %v", expected, result) + } + + httpHeader.Set("Content-Encoding", "gzip") + if result, expected := gzipReceived(httpHeader), true; result != expected { + t.Errorf("Expected %v but got: %v", expected, result) + } + + httpHeader.Set("Content-Encoding", "gzip, deflate, br") + if result, expected := gzipReceived(httpHeader), true; result != expected { + t.Errorf("Expected %v but got: %v", expected, result) + } + + httpHeader.Set("Content-Encoding", "br;q=1.0, gzip;q=0.8, *;q=0.1") + if result, expected := gzipReceived(httpHeader), true; result != expected { + t.Errorf("Expected %v but got: %v", expected, result) + } +} + func TestValidatePprofUrl(t *testing.T) { req := http.Request{} @@ -78,12 +106,15 @@ func TestRequestLogging(t *testing.T) { logger := test.New() logger.SetLevel(logging.Debug) + // set the threshold to a low value so the server compresses the response when asked to + gzipMinLength := "server.encoding.gzip.min_length=5" shutdownSeconds := 1 params := NewParams() params.Addrs = &[]string{":0"} params.Logger = logger params.PprofEnabled = true params.GracefulShutdownPeriod = shutdownSeconds // arbitrary, must be non-zero + params.ConfigOverrides = []string{gzipMinLength} rt, err := NewRuntime(ctx, params) if err != nil { @@ -98,36 +129,119 @@ func TestRequestLogging(t *testing.T) { }() <-initChannel + // prepare the request bodies to be used + var dataEndpointBody = []byte(`{"input": {"data": "checkForMe"}}`) + var compileEndpointBody = []byte(`{"unknowns": ["input"], "query": "data.checkForMe = true"}`) + var dataEndpointCompressedBody = zipString(`{"input": {"data": "checkForMe"}}`) + var compileEndpointCompressedBody = zipString(`{"unknowns": ["input"], "query": "data.checkForMe = true"}`) + tests := []struct { - path string - acceptEncoding string - expected string + path string + acceptEncoding string + expected string + expectedEncoding string + contentEncoding string + requestBody *[]byte }{ { - "/metrics", "gzip", "[compressed payload]", + path: "/metrics", + acceptEncoding: "gzip", + expected: "[compressed payload]", + expectedEncoding: "gzip", + contentEncoding: "", + requestBody: nil, }, { - "/metrics", "*/*", "HELP go_gc_duration_seconds A summary of the pause duration of garbage collection cycles.", // rest omitted + path: "/metrics", + acceptEncoding: "*/*", + expected: "HELP go_gc_duration_seconds A summary of the pause duration of garbage collection cycles.", + expectedEncoding: "", + contentEncoding: "", + requestBody: nil, }, - { // accept-encoding does not matter for "our" handlers -- they don't compress - "/v1/data", "gzip", "{\"result\":{}}", + { // the data handler on GET will compress the response if response is above server.encoding.gzip.min_length in size + path: "/v1/data", + acceptEncoding: "gzip", + expected: "{\"result\":{}}", + expectedEncoding: "gzip", + contentEncoding: "", + requestBody: nil, + }, + { // the data handler on POST can compress the response if response is above server.encoding.gzip.min_length in size + path: "/v1/data", + acceptEncoding: "gzip", + expected: "{\"result\":{}}", + expectedEncoding: "gzip", + contentEncoding: "", + requestBody: &dataEndpointBody, + }, + { // the data handler on POST can consume compressed request + path: "/v1/data", + acceptEncoding: "gzip", + expected: "{\"result\":{}}", + expectedEncoding: "gzip", + contentEncoding: "gzip", + requestBody: &dataEndpointCompressedBody, + }, + { // the compile handler will compress the response if response is above server.encoding.gzip.min_length in size + path: "/v1/compile", + acceptEncoding: "gzip", + expected: "{\"result\":{}}", + expectedEncoding: "gzip", + contentEncoding: "", + requestBody: &compileEndpointBody, + }, + { // the compile handler can consume compressed request + path: "/v1/compile", + acceptEncoding: "gzip", + expected: "{\"result\":{}}", + expectedEncoding: "gzip", + contentEncoding: "gzip", + requestBody: &compileEndpointCompressedBody, + }, + { // the handlers return plain data + path: "/v1/data", + acceptEncoding: "*/*", + expected: "{\"result\":{}}", + expectedEncoding: "", + contentEncoding: "", + requestBody: nil, }, { // accept-encoding does not matter for pprof: it's always protobuf - "/debug/pprof/cmdline", "*/*", "[binary payload]", + path: "/debug/pprof/cmdline", + acceptEncoding: "*/*", + expected: "[binary payload]", + expectedEncoding: "", + contentEncoding: "", + requestBody: nil, }, } // execute all the requests for _, tc := range tests { rec := httptest.NewRecorder() - req, err := http.NewRequest("GET", tc.path, nil) + method := "GET" + var body io.Reader + if tc.requestBody != nil { + method = "POST" + body = bytes.NewReader(*tc.requestBody) + } + req, err := http.NewRequest(method, tc.path, body) + if err != nil { t.Fatal(err) } req.Header.Set("Accept-Encoding", tc.acceptEncoding) + if tc.contentEncoding != "" { + req.Header.Set("Content-Encoding", tc.contentEncoding) + } rt.server.Handler.ServeHTTP(rec, req) if exp, act := http.StatusOK, rec.Result().StatusCode; exp != act { - t.Errorf("GET %s: expected HTTP %d, got %d", tc.path, exp, act) + t.Errorf("%s %s: expected HTTP %d, got %d", method, tc.path, exp, act) + } + contentEncoding := rec.Result().Header.Get("Content-Encoding") + if contentEncoding != tc.expectedEncoding { + t.Errorf("%s %s: expected content encoding %s, got %s", method, tc.path, tc.expectedEncoding, contentEncoding) } } @@ -137,19 +251,33 @@ func TestRequestLogging(t *testing.T) { ents := logger.Entries() for j, tc := range tests { i := uint64(j + 1) - found := false + foundResponse := false + foundRequest := false for _, ent := range entriesForReq(ents, i) { if ent.Message == "Sent response." { act := ent.Fields["resp_body"].(string) if !strings.Contains(act, tc.expected) { t.Errorf("expected %q in resp_body field, got %q", tc.expected, act) } - found = true + foundResponse = true + } + if tc.requestBody != nil && ent.Message == "Received request." { + if tc.requestBody != nil { + // the req_body is always uncompressed + act := ent.Fields["req_body"].(string) + if !strings.Contains(act, "checkForMe") { + t.Errorf("expected string %q in req_body field, got %q", "checkForMe", act) + } + foundRequest = true + } } } - if !found { + if !foundResponse { t.Errorf("Expected \"Sent response.\" log for request %d (path %s)", j, tc.path) } + if tc.requestBody != nil && !foundRequest { + t.Errorf("Expected \"Received request.\" log for request %d (path %s)", j, tc.path) + } } if t.Failed() { @@ -170,3 +298,14 @@ func entriesForReq(ents []test.LogEntry, n uint64) []test.LogEntry { } return ret } +func zipString(input string) []byte { + var b bytes.Buffer + gz := gzip.NewWriter(&b) + if _, err := gz.Write([]byte(input)); err != nil { + log.Fatal(err) + } + if err := gz.Close(); err != nil { + log.Fatal(err) + } + return b.Bytes() +} diff --git a/server/handlers/compress.go b/server/handlers/compress.go new file mode 100644 index 0000000000..1994269d10 --- /dev/null +++ b/server/handlers/compress.go @@ -0,0 +1,194 @@ +package handlers + +import ( + "compress/gzip" + "fmt" + "io" + "net/http" + "strings" + "sync" +) + +const ( + acceptEncodingHeader = "Accept-Encoding" + contentEncodingHeader = "Content-Encoding" + contentLengthHeader = "Content-Length" + gzipEncodingValue = "gzip" +) + +// This handler applies only for data and compile endpoints, for selected HTTP methods +// +// If the client asked for a gzip response, this handler will buffer the response and +// wait until it reached a certain threshold. If the threshold is not hit, the uncompressed response is sent +// +// If a gzip response is not asked by the client, it'll send the uncompressed response +// +// The threshold and the gzip compression level can be modified from server's configuration + +func CompressHandler(handler http.Handler, gzipMinLength int, gzipCompressionLevel int) http.Handler { + initGzipPool(gzipCompressionLevel) + + return http.HandlerFunc(func(responseWriter http.ResponseWriter, request *http.Request) { + enabledForEndpoint := isDataEndpoint(request) || isCompileEndpoint(request) + if !enabledForEndpoint { + handler.ServeHTTP(responseWriter, request) + return + } + + responseWriter.Header().Add("Vary", acceptEncodingHeader) + + if !gzipHeaderDetected(request.Header) { + handler.ServeHTTP(responseWriter, request) + return + } + + crw := &compressResponseWriter{ + ResponseWriter: responseWriter, + headerWritten: false, + minlength: gzipMinLength, + } + defer crw.Close() + handler.ServeHTTP(crw, request) + }) +} + +type compressResponseWriter struct { + gzipWriter *gzip.Writer + http.ResponseWriter + buffer []byte + statusCode int + headerWritten bool + minlength int +} + +var gzipPool *sync.Pool + +func initGzipPool(compressionLevel int) { + if gzipPool == nil { + gzipPool = &sync.Pool{ + New: func() interface{} { + writer, _ := gzip.NewWriterLevel(io.Discard, compressionLevel) + return writer + }, + } + } +} + +func (w *compressResponseWriter) WriteHeader(statusCode int) { + // save the status code for later use + w.statusCode = statusCode +} + +func (w *compressResponseWriter) Write(bytes []byte) (int, error) { + if w.isGzipInitialized() { + return w.gzipWriter.Write(bytes) + } + + // accumulate the buffer + w.buffer = append(w.buffer, bytes...) + + // if the buffer is above threshold, use compression + if len(w.buffer) >= w.minlength { + err := w.doCompressedResponse() + if err != nil { + return 0, err + } + return len(bytes), nil + } + + // wait for more data + return len(bytes), nil +} + +func (w *compressResponseWriter) Flush() { + if w.isGzipInitialized() { + w.gzipWriter.Flush() + flusher, canFlush := w.ResponseWriter.(http.Flusher) + if canFlush { + flusher.Flush() + } + } +} + +func (w *compressResponseWriter) Close() error { + if !w.isGzipInitialized() { + // gzip didn't handle the response, send it plain + err := w.doUncompressedResponse() + if err != nil { + err = fmt.Errorf("error writing uncompressed data: %v", err.Error()) + } + return err + } + + err := w.gzipWriter.Close() + defer gzipPool.Put(w.gzipWriter) + w.gzipWriter = nil + return err +} + +func (w *compressResponseWriter) doCompressedResponse() error { + w.ResponseWriter.Header().Set(contentEncodingHeader, gzipEncodingValue) + w.Header().Del(contentLengthHeader) + w.writeHeader() + // there's nothing to write + if w.buffer == nil || len(w.buffer) <= 0 { + return nil + } + gzipWriter := gzipPool.Get().(*gzip.Writer) + gzipWriter.Reset(w.ResponseWriter) + w.gzipWriter = gzipWriter + _, err := w.gzipWriter.Write(w.buffer) + return err +} + +func (w *compressResponseWriter) doUncompressedResponse() error { + w.writeHeader() + // there's nothing to write + if w.buffer == nil { + return nil + } + _, err := w.ResponseWriter.Write(w.buffer) + w.buffer = nil + return err +} + +func (w *compressResponseWriter) isGzipInitialized() bool { + return w.gzipWriter != nil +} + +func (w *compressResponseWriter) writeHeader() { + if !w.headerWritten && w.statusCode != 0 { + w.ResponseWriter.WriteHeader(w.statusCode) + w.headerWritten = true + } +} + +func isDataEndpoint(req *http.Request) bool { + isPostOrGetMethod := isPostMethod(req) || isGetMethod(req) + isV1rV0 := strings.HasPrefix(req.URL.Path, "/v1/data") || strings.HasPrefix(req.URL.Path, "/v0/data") + return isPostOrGetMethod && isV1rV0 +} + +func isCompileEndpoint(req *http.Request) bool { + return isPostMethod(req) && strings.HasPrefix(req.URL.Path, "/v1/compile") +} + +func isPostMethod(req *http.Request) bool { + return req.Method == "POST" +} + +func isGetMethod(req *http.Request) bool { + return req.Method == "GET" +} + +func gzipHeaderDetected(header http.Header) bool { + a := header.Get("Accept-Encoding") + parts := strings.Split(a, ",") + for _, part := range parts { + part = strings.TrimSpace(part) + if part == gzipEncodingValue || strings.HasPrefix(part, gzipEncodingValue+";") { + return true + } + } + return false +} diff --git a/server/handlers/compress_test.go b/server/handlers/compress_test.go new file mode 100644 index 0000000000..1da6649726 --- /dev/null +++ b/server/handlers/compress_test.go @@ -0,0 +1,182 @@ +package handlers + +import ( + "bytes" + "compress/gzip" + "io" + "log" + "net/http" + "net/http/httptest" + "net/url" + "testing" +) + +const ( + gzipEncoding = "gzip" + requestBody = "Hello World!\n" +) + +var defaultCompressionLevel = gzip.BestCompression + +type compressHandlerTestScenario struct { + path string + method string + acceptEncoding string + gzipMinSize int + expectedCompressedResponse bool +} + +func executeRequest(w *httptest.ResponseRecorder, testScenario compressHandlerTestScenario) { + CompressHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, err := io.WriteString(w, requestBody) + if err != nil { + log.Fatalf("Error writing the request body: %v", err) + } + }), testScenario.gzipMinSize, defaultCompressionLevel).ServeHTTP(w, &http.Request{ + URL: &url.URL{Path: testScenario.path}, + Method: testScenario.method, + Header: http.Header{ + "Accept-Encoding": []string{testScenario.acceptEncoding}, + }, + }) +} + +func TestCompressHandlerWithGzipOnInScopeEndpoints(t *testing.T) { + tests := map[string]compressHandlerTestScenario{ + "v0PostDataCompressed": { + path: "/v0/data", + method: "POST", + acceptEncoding: gzipEncoding, + gzipMinSize: 1, + expectedCompressedResponse: true, + }, + "v1PostDataCompressed": { + path: "/v1/data", + method: "POST", + acceptEncoding: gzipEncoding, + gzipMinSize: 1, + expectedCompressedResponse: true, + }, + "v1PostCompileCompressed": { + path: "/v1/compile", + method: "POST", + acceptEncoding: gzipEncoding, + gzipMinSize: 1, + expectedCompressedResponse: true, + }, + "v0PostDataUncompressed": { + path: "/v0/data", + method: "POST", + acceptEncoding: gzipEncoding, + gzipMinSize: 1024, + expectedCompressedResponse: false, + }, + "v1PostDataUncompressed": { + path: "/v1/data", + method: "POST", + acceptEncoding: gzipEncoding, + gzipMinSize: 1024, + expectedCompressedResponse: false, + }, + "v1PostCompileUncompressed": { + path: "/v1/compile", + method: "POST", + acceptEncoding: gzipEncoding, + gzipMinSize: 1024, + expectedCompressedResponse: false, + }, + "v1PostCompileAcceptEncodingAll": { + path: "/v1/compile", + method: "POST", + acceptEncoding: "*/*", + gzipMinSize: 1024, + expectedCompressedResponse: false, + }, + } + for name, ts := range tests { + w := httptest.NewRecorder() + executeRequest(w, ts) + if w.Result().Header.Get("Vary") != "Accept-Encoding" { + t.Error("missing the Vary header") + } + contentEncodingValue := w.Result().Header.Get("Content-Encoding") + if ts.expectedCompressedResponse { + if contentEncodingValue != gzipEncoding { + t.Errorf("wrong content encoding, got %q want %q", contentEncodingValue, gzipEncoding) + } + expectedLength := len(zipString(requestBody)) + receivedLength := w.Body.Len() + if receivedLength != expectedLength { + t.Errorf("test: %s wrong len, got %d want %d", name, w.Body.Len(), expectedLength) + } + receivedBody := unzip(w.Body.Bytes()) + if receivedBody != requestBody { + t.Errorf("test: %s wrong body, got %v, want %v", name, receivedBody, requestBody) + } + } else { + if contentEncodingValue == gzipEncoding { + t.Errorf("wrong content encoding, got %q want %q", contentEncodingValue, "") + } + expectedLength := len(requestBody) + receivedLength := w.Body.Len() + if receivedLength != expectedLength { + t.Errorf("test: %s wrong len, got %d want %d", name, w.Body.Len(), expectedLength) + } + receivedBody := w.Body.String() + if receivedBody != requestBody { + t.Errorf("test: %s wrong body, got %v, want %v", name, receivedBody, requestBody) + } + } + } +} + +func TestHandlerOnEndpointsWithoutCompression(t *testing.T) { + testScenario := compressHandlerTestScenario{ + path: "/metrics", + method: "GET", + acceptEncoding: gzipEncoding, + gzipMinSize: 1, + } + w := httptest.NewRecorder() + executeRequest(w, testScenario) + contentEncodingValue := w.Result().Header.Get("Content-Encoding") + if contentEncodingValue != "" { + t.Errorf("wrong content encoding, got %q want %q", contentEncodingValue, gzipEncoding) + } + + expectedLength := len(requestBody) + receivedLength := w.Body.Len() + if receivedLength != expectedLength { + t.Errorf("wrong len, got %d want %d", w.Body.Len(), expectedLength) + } +} + +func zipString(input string) []byte { + var b bytes.Buffer + gz := gzip.NewWriter(&b) + if _, err := gz.Write([]byte(input)); err != nil { + log.Fatal(err) + } + if err := gz.Close(); err != nil { + log.Fatal(err) + } + return b.Bytes() +} + +func unzip(body []byte) string { + reader := bytes.NewReader(body) + gzReader, err := gzip.NewReader(reader) + if err != nil { + log.Fatalf("Unexpected gzip error: %v", err) + } + plainOutput, err := io.ReadAll(gzReader) + if err != nil { + log.Fatalf("Unexpected gzip error: %v", err) + } + err = gzReader.Close() + if err != nil { + log.Fatalf("Unexpected gzip close err: %v", err) + } + return string(plainOutput) +} diff --git a/server/server.go b/server/server.go index 3c1d4046ed..7ea9c3398d 100644 --- a/server/server.go +++ b/server/server.go @@ -6,6 +6,7 @@ package server import ( "bytes" + "compress/gzip" "context" "crypto/tls" "crypto/x509" @@ -24,6 +25,8 @@ import ( "sync" "time" + serverEncodingPlugin "github.com/open-policy-agent/opa/plugins/server/encoding" + "github.com/gorilla/mux" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" @@ -41,6 +44,7 @@ import ( "github.com/open-policy-agent/opa/plugins/status" "github.com/open-policy-agent/opa/rego" "github.com/open-policy-agent/opa/server/authorizer" + "github.com/open-policy-agent/opa/server/handlers" "github.com/open-policy-agent/opa/server/identifier" "github.com/open-policy-agent/opa/server/types" "github.com/open-policy-agent/opa/server/writer" @@ -185,6 +189,11 @@ func (s *Server) Init(ctx context.Context) (*Server, error) { // authorizer, if configured, needs the iCache to be set up already s.Handler = s.initHandlerAuth(s.Handler) + // compression handler + s.Handler, err = s.initHandlerCompression(s.Handler) + if err != nil { + return nil, err + } s.DiagnosticHandler = s.initHandlerAuth(s.DiagnosticHandler) return s, s.store.Commit(ctx, txn) @@ -658,6 +667,21 @@ func (s *Server) initHandlerAuth(handler http.Handler) http.Handler { return handler } +func (s *Server) initHandlerCompression(handler http.Handler) (http.Handler, error) { + var encodingRawConfig json.RawMessage + serverConfig := s.manager.Config.Server + if serverConfig != nil { + encodingRawConfig = serverConfig.Encoding + } + encodingConfig, err := serverEncodingPlugin.NewConfigBuilder().WithBytes(encodingRawConfig).Parse() + if err != nil { + return nil, err + } + compressHandler := handlers.CompressHandler(handler, *encodingConfig.Gzip.MinLength, *encodingConfig.Gzip.CompressionLevel) + + return compressHandler, nil +} + func (s *Server) initRouters() { mainRouter := s.router if mainRouter == nil { @@ -1243,7 +1267,15 @@ func (s *Server) v1CompilePost(w http.ResponseWriter, r *http.Request) { m.Timer(metrics.ServerHandler).Start() m.Timer(metrics.RegoQueryParse).Start() - request, reqErr := readInputCompilePostV1(r.Body) + // decompress the input if sent as zip + body, err := readPlainBody(r) + if err != nil { + reqErr := types.NewErrorV1(types.CodeInvalidParameter, "could not decompress the body") + writer.Error(w, http.StatusBadRequest, reqErr) + return + } + + request, reqErr := readInputCompilePostV1(body) if reqErr != nil { writer.Error(w, http.StatusBadRequest, reqErr) return @@ -2713,10 +2745,16 @@ func readInputV0(r *http.Request) (ast.Value, error) { return ast.InterfaceToValue(parsed) } + // decompress the input if sent as zip + body, err := readPlainBody(r) + if err != nil { + return nil, fmt.Errorf("could not decompress the body: %w", err) + } + var x interface{} if strings.Contains(r.Header.Get("Content-Type"), "yaml") { - bs, err := io.ReadAll(r.Body) + bs, err := io.ReadAll(body) if err != nil { return nil, err } @@ -2726,7 +2764,7 @@ func readInputV0(r *http.Request) (ast.Value, error) { } } } else { - dec := util.NewJSONDecoder(r.Body) + dec := util.NewJSONDecoder(body) if err := dec.Decode(&x); err != nil && err != io.EOF { return nil, fmt.Errorf("body contains malformed input document: %w", err) } @@ -2757,11 +2795,17 @@ func readInputPostV1(r *http.Request) (ast.Value, error) { var request types.DataRequestV1 + // decompress the input if sent as zip + body, err := readPlainBody(r) + if err != nil { + return nil, fmt.Errorf("could not decompress the body: %w", err) + } + ct := r.Header.Get("Content-Type") // There is no standard for yaml mime-type so we just look for // anything related if strings.Contains(ct, "yaml") { - bs, err := io.ReadAll(r.Body) + bs, err := io.ReadAll(body) if err != nil { return nil, err } @@ -2771,7 +2815,7 @@ func readInputPostV1(r *http.Request) (ast.Value, error) { } } } else { - dec := util.NewJSONDecoder(r.Body) + dec := util.NewJSONDecoder(body) if err := dec.Decode(&request); err != nil && err != io.EOF { return nil, fmt.Errorf("body contains malformed input document: %w", err) } @@ -2981,3 +3025,19 @@ func annotateSpan(ctx context.Context, decisionID string) { trace.SpanFromContext(ctx). SetAttributes(attribute.String(otelDecisionIDAttr, decisionID)) } + +func readPlainBody(r *http.Request) (io.ReadCloser, error) { + if strings.Contains(r.Header.Get("Content-Encoding"), "gzip") { + gzReader, err := gzip.NewReader(r.Body) + if err != nil { + return nil, err + } + bytesBody, err := io.ReadAll(gzReader) + if err != nil { + return nil, err + } + defer gzReader.Close() + return io.NopCloser(bytes.NewReader(bytesBody)), err + } + return r.Body, nil +} diff --git a/server/server_test.go b/server/server_test.go index 5803334f4e..a9bf0391b4 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -7,10 +7,13 @@ package server import ( "bytes" + "compress/gzip" "context" "encoding/json" "errors" "fmt" + "io" + "log" "net/http" "net/http/httptest" "net/url" @@ -1763,6 +1766,343 @@ func TestDataPutV1IfNoneMatch(t *testing.T) { } } +func TestDataPostV0CompressedResponse(t *testing.T) { + tests := []struct { + gzipMinLength int + compressedResponse bool + }{ + { + gzipMinLength: 3, + compressedResponse: true, + }, + { + gzipMinLength: 1400, + compressedResponse: false, + }, + } + + for _, test := range tests { + f := newFixtureWithConfig(t, fmt.Sprintf(`{"server":{"encoding":{"gzip":{"min_length": %d}}}}`, test.gzipMinLength)) + // create the policy + err := f.v1(http.MethodPut, "/policies/test", `package opa.examples +import input.example.flag +allow_request { flag == true } +`, 200, "") + if err != nil { + t.Fatal(err) + } + + // execute the request + req := newReqV0(http.MethodPost, "/data/opa/examples/allow_request", `{"example": {"flag": true}}`) + req.Header.Set("Accept-Encoding", "gzip") + f.reset() + f.server.Handler.ServeHTTP(f.recorder, req) + + // check for content encoding + expectedEncoding := "gzip" + if !test.compressedResponse { + expectedEncoding = "" + } + receivedEncodingHeaderValue := f.recorder.Header().Get("Content-Encoding") + if receivedEncodingHeaderValue != expectedEncoding { + t.Fatalf("Expected Content-Encoding %v but got: %v", expectedEncoding, receivedEncodingHeaderValue) + } + + var plainOutput []byte + if test.compressedResponse { + // unzip the response + gzReader, err := gzip.NewReader(f.recorder.Body) + if err != nil { + t.Fatalf("Unexpected gzip error: %v", err) + } + plainOutput, err = io.ReadAll(gzReader) + if err != nil { + t.Fatalf("Unexpected error on reading the response: %v", err) + } + } else { + plainOutput = f.recorder.Body.Bytes() + } + + expected := "true" + result := strings.TrimSuffix(string(plainOutput), "\n") + if plainOutput == nil || result != expected { + t.Fatalf("Expected %v but got: %v", expected, result) + } + } +} + +func TestDataPostV1CompressedResponse(t *testing.T) { + tests := []struct { + gzipMinLength int + compressedResponse bool + }{ + { + gzipMinLength: 3, + compressedResponse: true, + }, + { + gzipMinLength: 1400, + compressedResponse: false, + }, + } + + for _, test := range tests { + f := newFixtureWithConfig(t, fmt.Sprintf(`{"server":{"encoding":{"gzip":{"min_length": %d}}}}`, test.gzipMinLength)) + // create the policy + err := f.v1(http.MethodPut, "/policies/test", `package test +default hello := false +hello { + input.message == "world" +} +`, 200, "") + if err != nil { + t.Fatal(err) + } + + // execute the request + req := newReqV1(http.MethodPost, "/data/test", `{"input": {"message": "world"}}`) + req.Header.Set("Accept-Encoding", "gzip") + f.reset() + f.server.Handler.ServeHTTP(f.recorder, req) + + var result types.DataResponseV1 + + // check for content encoding + expectedEncoding := "gzip" + if !test.compressedResponse { + expectedEncoding = "" + } + receivedEncodingHeaderValue := f.recorder.Header().Get("Content-Encoding") + if receivedEncodingHeaderValue != expectedEncoding { + t.Fatalf("Expected Content-Encoding %v but got: %v", expectedEncoding, receivedEncodingHeaderValue) + } + + if test.compressedResponse { + // unzip and unmarshall the response + gzReader, err := gzip.NewReader(f.recorder.Body) + if err != nil { + t.Fatalf("Unexpected gzip error: %v", err) + } + if err := util.NewJSONDecoder(gzReader).Decode(&result); err != nil { + t.Fatalf("Unexpected JSON decode error: %v", err) + } + } else { + // unmarshall the response + if err := util.NewJSONDecoder(f.recorder.Body).Decode(&result); err != nil { + t.Fatalf("Unexpected JSON decode error: %v", err) + } + } + + var expected interface{} + if err := util.UnmarshalJSON([]byte(`{"hello": true}`), &expected); err != nil { + panic(err) + } + if result.Result == nil || !reflect.DeepEqual(*result.Result, expected) { + t.Fatalf("Expected %v but got: %v", expected, *result.Result) + } + } +} + +func TestCompileV1CompressedResponse(t *testing.T) { + tests := []struct { + gzipMinLength int + compressedResponse bool + }{ + { + gzipMinLength: 3, + compressedResponse: true, + }, + { + gzipMinLength: 1400, + compressedResponse: false, + }, + } + + for _, test := range tests { + f := newFixtureWithConfig(t, fmt.Sprintf(`{"server":{"encoding":{"gzip":{"min_length": %d}}}}`, test.gzipMinLength)) + + // create the policy + mod := `package test + + p { + input.x = 1 + } + + q { + data.a[i] = input.x + } + + default r = true + + r { input.x = 1 } + + custom_func(x) { data.a[i] == x } + + s { custom_func(input.x) } + ` + err := f.v1(http.MethodPut, "/policies/test", mod, 200, "") + if err != nil { + t.Fatal(err) + } + + // execute the request + req := newReqV1(http.MethodPost, "/compile", `{"unknowns": ["input"], "query": "data.test.p = true"}`) + req.Header.Set("Accept-Encoding", "gzip") + f.reset() + f.server.Handler.ServeHTTP(f.recorder, req) + + var result types.CompileResponseV1 + + // check for content encoding + expectedEncoding := "gzip" + if !test.compressedResponse { + expectedEncoding = "" + } + receivedEncodingHeaderValue := f.recorder.Header().Get("Content-Encoding") + if receivedEncodingHeaderValue != expectedEncoding { + t.Fatalf("Expected Content-Encoding %v but got: %v", expectedEncoding, receivedEncodingHeaderValue) + } + + if test.compressedResponse { + // unzip and unmarshall the response + gzReader, err := gzip.NewReader(f.recorder.Body) + if err != nil { + t.Fatalf("Unexpected gzip error: %v", err) + } + if err := util.NewJSONDecoder(gzReader).Decode(&result); err != nil { + t.Fatalf("Unexpected JSON decode error: %v", err) + } + } else { + // unmarshall the response + if err := util.NewJSONDecoder(f.recorder.Body).Decode(&result); err != nil { + t.Fatalf("Unexpected JSON decode error: %v", err) + } + } + + var expected interface{} + expectedStr := fmt.Sprintf(`{"queries": [%v]}`, string(util.MustMarshalJSON(ast.MustParseBody("input.x = 1")))) + if err := util.UnmarshalJSON([]byte(expectedStr), &expected); err != nil { + panic(err) + } + + if result.Result == nil || !reflect.DeepEqual(*result.Result, expected) { + t.Fatalf("Expected %v but got: %v", expected, *result.Result) + } + } +} + +func TestDataPostV0CompressedRequest(t *testing.T) { + f := newFixture(t) + // create the policy + err := f.v1(http.MethodPut, "/policies/test", `package opa.examples +import input.example.flag +allow_request { flag == true } +`, 200, "") + if err != nil { + t.Fatal(err) + } + + // execute the request + compressedBoy := zipString(`{"example": {"flag": true}}`) + req := newStreamedReqV0(http.MethodPost, "/data/opa/examples/allow_request", bytes.NewReader(compressedBoy)) + req.Header.Set("Content-Encoding", "gzip") + f.reset() + f.server.Handler.ServeHTTP(f.recorder, req) + + expected := "true" + result := strings.TrimSuffix(f.recorder.Body.String(), "\n") + if result != expected { + t.Fatalf("Expected %v but got: %v", expected, result) + } +} + +func TestDataPostV1CompressedRequest(t *testing.T) { + f := newFixture(t) + // create the policy + err := f.v1(http.MethodPut, "/policies/test", `package test +default hello := false +hello { + input.message == "world" +} +`, 200, "") + if err != nil { + t.Fatal(err) + } + + // execute the request + compressedBoy := zipString(`{"input": {"message": "world"}}`) + req := newStreamedReqV1(http.MethodPost, "/data/test", bytes.NewReader(compressedBoy)) + req.Header.Set("Content-Encoding", "gzip") + f.reset() + f.server.Handler.ServeHTTP(f.recorder, req) + + var result types.DataResponseV1 + + // unmarshall the response + if err := util.NewJSONDecoder(f.recorder.Body).Decode(&result); err != nil { + t.Fatalf("Unexpected JSON decode error: %v", err) + } + + var expected interface{} + if err := util.UnmarshalJSON([]byte(`{"hello": true}`), &expected); err != nil { + panic(err) + } + if result.Result == nil || !reflect.DeepEqual(*result.Result, expected) { + t.Fatalf("Expected %v but got: %v", expected, *result.Result) + } +} + +func TestCompileV1CompressedRequest(t *testing.T) { + f := newFixture(t) + // create the policy + mod := `package test + + p { + input.x = 1 + } + + q { + data.a[i] = input.x + } + + default r = true + + r { input.x = 1 } + + custom_func(x) { data.a[i] == x } + + s { custom_func(input.x) } + ` + err := f.v1(http.MethodPut, "/policies/test", mod, 200, "") + if err != nil { + t.Fatal(err) + } + + // execute the request + compressedBoy := zipString(`{"unknowns": ["input"], "query": "data.test.p = true"}`) + req := newStreamedReqV1(http.MethodPost, "/compile", bytes.NewReader(compressedBoy)) + req.Header.Set("Content-Encoding", "gzip") + f.reset() + f.server.Handler.ServeHTTP(f.recorder, req) + + var result types.CompileResponseV1 + + // unmarshall the response + if err := util.NewJSONDecoder(f.recorder.Body).Decode(&result); err != nil { + t.Fatalf("Unexpected JSON decode error: %v", err) + } + + var expected interface{} + expectedStr := fmt.Sprintf(`{"queries": [%v]}`, string(util.MustMarshalJSON(ast.MustParseBody("input.x = 1")))) + if err := util.UnmarshalJSON([]byte(expectedStr), &expected); err != nil { + panic(err) + } + + if result.Result == nil || !reflect.DeepEqual(*result.Result, expected) { + t.Fatalf("Expected %v but got: %v", expected, *result.Result) + } +} + func TestBundleScope(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) @@ -3926,6 +4266,36 @@ func newFixture(t *testing.T, opts ...func(*Server)) *fixture { } } +func newFixtureWithConfig(t *testing.T, config string, opts ...func(*Server)) *fixture { + ctx := context.Background() + server := New(). + WithAddresses([]string{"localhost:8182"}). + WithStore(inmem.New()) // potentially overridden via opts + for _, opt := range opts { + opt(server) + } + + m, err := plugins.New([]byte(config), "test", server.store) + if err != nil { + t.Fatal(err) + } + server = server.WithManager(m) + if err := m.Start(ctx); err != nil { + t.Fatal(err) + } + server, err = server.Init(ctx) + if err != nil { + t.Fatal(err) + } + recorder := httptest.NewRecorder() + + return &fixture{ + server: server, + recorder: recorder, + t: t, + } +} + func newFixtureWithStore(t *testing.T, store storage.Store, opts ...func(*Server)) *fixture { ctx := context.Background() m, err := plugins.New([]byte{}, "test", store) @@ -4097,6 +4467,26 @@ func newReqUnversioned(method, path, body string) *http.Request { return req } +func newStreamedReqV0(method string, path string, body io.Reader) *http.Request { + return newStreamedReq(0, method, path, body) +} + +func newStreamedReqV1(method string, path string, body io.Reader) *http.Request { + return newStreamedReq(1, method, path, body) +} + +func newStreamedReq(version int, method string, path string, body io.Reader) *http.Request { + return newStreamedReqUnversioned(method, fmt.Sprintf("/v%d", version)+path, body) +} + +func newStreamedReqUnversioned(method string, path string, body io.Reader) *http.Request { + req, err := http.NewRequest(method, path, body) + if err != nil { + panic(err) + } + return req +} + func mustUnmarshalTrace(t types.TraceV1) (trace types.TraceV1Raw) { if err := json.Unmarshal(t, &trace); err != nil { panic("not reached") @@ -4478,3 +4868,15 @@ func (m mockHTTPListener) Shutdown(context.Context) error { func (m mockHTTPListener) Type() httpListenerType { return m.t } + +func zipString(input string) []byte { + var b bytes.Buffer + gz := gzip.NewWriter(&b) + if _, err := gz.Write([]byte(input)); err != nil { + log.Fatal(err) + } + if err := gz.Close(); err != nil { + log.Fatal(err) + } + return b.Bytes() +}