Created common methods for keep and map poll responses

This commit is contained in:
Juan Font Alonso 2022-08-14 22:50:39 +02:00
parent 35f3dee1d0
commit f4bab6b290

View File

@ -10,35 +10,87 @@ import (
"tailscale.com/types/key"
)
func (h *Headscale) getMapKeepAliveResponse(
machineKey key.MachinePublic,
func (h *Headscale) getMapResponseData(
mapRequest tailcfg.MapRequest,
machine *Machine,
isNoise bool,
) ([]byte, error) {
mapResponse := tailcfg.MapResponse{
mapResponse, err := h.generateMapResponse(mapRequest, machine)
if err != nil {
return nil, err
}
if isNoise {
return h.marshalResponse(mapResponse, mapRequest.Compress, key.MachinePublic{})
}
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.marshalResponse(mapResponse, mapRequest.Compress, machineKey)
}
func (h *Headscale) getMapKeepAliveResponseData(
mapRequest tailcfg.MapRequest,
machine *Machine,
isNoise bool,
) ([]byte, error) {
keepAliveResponse := tailcfg.MapResponse{
KeepAlive: true,
}
var respBody []byte
var err error
if mapRequest.Compress == ZstdCompression {
src, err := json.Marshal(mapResponse)
if err != nil {
log.Error().
Caller().
Str("func", "getMapKeepAliveResponse").
Err(err).
Msg("Failed to marshal keepalive response for the client")
return nil, err
}
if isNoise {
return h.marshalResponse(keepAliveResponse, mapRequest.Compress, key.MachinePublic{})
}
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.marshalResponse(keepAliveResponse, mapRequest.Compress, machineKey)
}
func (h *Headscale) marshalResponse(
resp interface{},
compression string,
machineKey key.MachinePublic,
) ([]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)
srcCompressed := encoder.EncodeAll(src, nil)
respBody = h.privateKey.SealTo(machineKey, srcCompressed)
respBody = encoder.EncodeAll(jsonBody, nil)
if !machineKey.IsZero() { // if legacy protocol
respBody = h.privateKey.SealTo(machineKey, respBody)
}
} else {
respBody, err = encode(mapResponse, &machineKey, h.privateKey)
if err != nil {
return nil, err
if !machineKey.IsZero() { // if legacy protocol
respBody = h.privateKey.SealTo(machineKey, jsonBody)
}
}
data := make([]byte, reservedResponseHeaderSize)
binary.LittleEndian.PutUint32(data, uint32(len(respBody)))
data = append(data, respBody...)