1
0
mirror of https://github.com/juanfont/headscale.git synced 2024-12-20 19:09:07 +01:00
juanfont.headscale/protocol_common_utils.go

125 lines
2.6 KiB
Go
Raw Normal View History

package headscale
import (
"encoding/binary"
"encoding/json"
"github.com/klauspost/compress/zstd"
"github.com/rs/zerolog/log"
"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")
2022-09-04 11:36:03 +02:00
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 {
encoder, _ := zstd.NewWriter(nil)
respBody = encoder.EncodeAll(jsonBody, nil)
if !isNoise { // if legacy protocol
respBody = h.privateKey.SealTo(machineKey, respBody)
}
} else {
if !isNoise { // if legacy protocol
respBody = h.privateKey.SealTo(machineKey, jsonBody)
2022-08-14 23:15:41 +02:00
} else {
respBody = jsonBody
}
}
data := make([]byte, reservedResponseHeaderSize)
binary.LittleEndian.PutUint32(data, uint32(len(respBody)))
data = append(data, respBody...)
return data, nil
}