mirror of
https://github.com/juanfont/headscale.git
synced 2024-12-20 19:09:07 +01:00
4a7921ead5
Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
151 lines
3.1 KiB
Go
151 lines
3.1 KiB
Go
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"
|
|
)
|
|
|
|
func (h *Headscale) getMapResponseData(
|
|
mapRequest tailcfg.MapRequest,
|
|
machine *Machine,
|
|
isNoise bool,
|
|
) ([]byte, error) {
|
|
mapResponse, err := h.generateMapResponse(mapRequest, machine)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if isNoise {
|
|
return h.marshalMapResponse(mapResponse, key.MachinePublic{}, mapRequest.Compress, isNoise)
|
|
}
|
|
|
|
var machineKey key.MachinePublic
|
|
err = machineKey.UnmarshalText([]byte(MachinePublicKeyEnsurePrefix(machine.MachineKey)))
|
|
if err != nil {
|
|
log.Error().
|
|
Caller().
|
|
Err(err).
|
|
Msg("Cannot parse client key")
|
|
|
|
return nil, err
|
|
}
|
|
|
|
return h.marshalMapResponse(mapResponse, machineKey, mapRequest.Compress, isNoise)
|
|
}
|
|
|
|
func (h *Headscale) getMapKeepAliveResponseData(
|
|
mapRequest tailcfg.MapRequest,
|
|
machine *Machine,
|
|
isNoise bool,
|
|
) ([]byte, error) {
|
|
keepAliveResponse := tailcfg.MapResponse{
|
|
KeepAlive: true,
|
|
}
|
|
|
|
if isNoise {
|
|
return h.marshalMapResponse(keepAliveResponse, key.MachinePublic{}, mapRequest.Compress, isNoise)
|
|
}
|
|
|
|
var machineKey key.MachinePublic
|
|
err := machineKey.UnmarshalText([]byte(MachinePublicKeyEnsurePrefix(machine.MachineKey)))
|
|
if err != nil {
|
|
log.Error().
|
|
Caller().
|
|
Err(err).
|
|
Msg("Cannot parse client key")
|
|
|
|
return nil, err
|
|
}
|
|
|
|
return h.marshalMapResponse(keepAliveResponse, machineKey, mapRequest.Compress, isNoise)
|
|
}
|
|
|
|
func (h *Headscale) marshalResponse(
|
|
resp interface{},
|
|
machineKey key.MachinePublic,
|
|
isNoise bool,
|
|
) ([]byte, error) {
|
|
jsonBody, err := json.Marshal(resp)
|
|
if err != nil {
|
|
log.Error().
|
|
Caller().
|
|
Err(err).
|
|
Msg("Cannot marshal response")
|
|
|
|
return nil, err
|
|
}
|
|
|
|
if isNoise {
|
|
return jsonBody, nil
|
|
}
|
|
|
|
return h.privateKey.SealTo(machineKey, jsonBody), nil
|
|
}
|
|
|
|
func (h *Headscale) marshalMapResponse(
|
|
resp interface{},
|
|
machineKey key.MachinePublic,
|
|
compression string,
|
|
isNoise bool,
|
|
) ([]byte, error) {
|
|
jsonBody, err := json.Marshal(resp)
|
|
if err != nil {
|
|
log.Error().
|
|
Caller().
|
|
Err(err).
|
|
Msg("Cannot marshal map response")
|
|
}
|
|
|
|
var respBody []byte
|
|
if compression == ZstdCompression {
|
|
respBody = zstdEncode(jsonBody)
|
|
if !isNoise { // if legacy protocol
|
|
respBody = h.privateKey.SealTo(machineKey, respBody)
|
|
}
|
|
} else {
|
|
if !isNoise { // if legacy protocol
|
|
respBody = h.privateKey.SealTo(machineKey, jsonBody)
|
|
} else {
|
|
respBody = jsonBody
|
|
}
|
|
}
|
|
|
|
data := make([]byte, reservedResponseHeaderSize)
|
|
binary.LittleEndian.PutUint32(data, uint32(len(respBody)))
|
|
data = append(data, respBody...)
|
|
|
|
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
|
|
},
|
|
}
|