diff --git a/flake.nix b/flake.nix index 0167797..8a3e59f 100644 --- a/flake.nix +++ b/flake.nix @@ -33,7 +33,7 @@ # When updating go.mod or go.sum, a new sha will need to be calculated, # update this if you have a mismatch after doing a change to thos files. - vendorSha256 = "sha256-SuKT+b8g6xEK15ry2IAmpS/vwDG+zJqK9nfsWpHNXuU="; + vendorSha256 = "sha256-8p5NFxXKaZPsW4B6NMzfi0pqfVroIahSgA0fukvB3JI="; ldflags = ["-s" "-w" "-X github.com/juanfont/headscale/cmd/headscale/cli.Version=v${version}"]; }; diff --git a/go.mod b/go.mod index cba5e99..f35f19c 100644 --- a/go.mod +++ b/go.mod @@ -7,6 +7,7 @@ require ( github.com/ccding/go-stun/stun v0.0.0-20200514191101-4dc67bcdb029 github.com/cenkalti/backoff/v4 v4.2.0 github.com/coreos/go-oidc/v3 v3.4.0 + github.com/davecgh/go-spew v1.1.1 github.com/deckarep/golang-set/v2 v2.1.0 github.com/efekarakus/termcolor v1.0.1 github.com/glebarez/sqlite v1.5.0 @@ -58,7 +59,6 @@ require ( github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/containerd/console v1.0.3 // indirect github.com/containerd/continuity v0.3.0 // indirect - github.com/davecgh/go-spew v1.1.1 // indirect github.com/docker/cli v20.10.21+incompatible // indirect github.com/docker/docker v20.10.21+incompatible // indirect github.com/docker/go-connections v0.4.0 // indirect diff --git a/protocol_common_utils.go b/protocol_common_utils.go index d3fb697..96c236d 100644 --- a/protocol_common_utils.go +++ b/protocol_common_utils.go @@ -3,9 +3,11 @@ package headscale import ( "encoding/binary" "encoding/json" + "sync" "github.com/klauspost/compress/zstd" "github.com/rs/zerolog/log" + "tailscale.com/smallzstd" "tailscale.com/tailcfg" "tailscale.com/types/key" ) @@ -103,8 +105,7 @@ func (h *Headscale) marshalMapResponse( var respBody []byte if compression == ZstdCompression { - encoder, _ := zstd.NewWriter(nil) - respBody = encoder.EncodeAll(jsonBody, nil) + respBody = zstdEncode(jsonBody) if !isNoise { // if legacy protocol respBody = h.privateKey.SealTo(machineKey, respBody) } @@ -122,3 +123,28 @@ func (h *Headscale) marshalMapResponse( return data, nil } + +func zstdEncode(in []byte) []byte { + encoder, ok := zstdEncoderPool.Get().(*zstd.Encoder) + if !ok { + panic("invalid type in sync pool") + } + out := encoder.EncodeAll(in, nil) + _ = encoder.Close() + zstdEncoderPool.Put(encoder) + + return out +} + +var zstdEncoderPool = &sync.Pool{ + New: func() any { + encoder, err := smallzstd.NewEncoder( + nil, + zstd.WithEncoderLevel(zstd.SpeedFastest)) + if err != nil { + panic(err) + } + + return encoder + }, +}