Something went wrong. Try again.
forked from tailscale/tailscale
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425// Copyright (c) Tailscale Inc & contributors// SPDX-License-Identifier: BSD-3-Clause
package tsp
import ( "bytes" "cmp" "context" "encoding/binary" "encoding/json" "errors" "fmt" "io" "net/http" "strings" "sync" "sync/atomic"
"github.com/klauspost/compress/zstd" "tailscale.com/control/ts2021" "tailscale.com/tailcfg" "tailscale.com/types/key" "tailscale.com/util/httpbody")
// errSessionClosed is returned by [MapSession.Next] and// [MapSession.NextInto] when called after [MapSession.Close].var errSessionClosed = errors.New("tsp: map session closed")
// DefaultMaxMessageSize is the default cap, in bytes, on the size of a// single compressed map response frame. See [MapOpts.MaxMessageSize].const DefaultMaxMessageSize = 4 << 20
// zstdDecoderPool is a pool of *zstd.Decoder reused across MapSessions to// amortize the cost of setting up zstd state. Decoders are returned via// [MapSession.Close]; entries are reclaimed by the runtime under memory// pressure via sync.Pool semantics.var zstdDecoderPool sync.Pool // of *zstd.Decoder
// MapOpts contains options for sending a map request.type MapOpts struct { // NodeKey is the node's private key. Required. NodeKey key.NodePrivate
// Hostinfo is the host information to send. Optional; // if nil, a minimal default is used. Hostinfo *tailcfg.Hostinfo
// Stream is whether to receive multiple MapResponses over // the same HTTP connection. Stream bool
// OmitPeers is whether the client is okay with the Peers list // being omitted in the response. OmitPeers bool
// MaxMessageSize is the maximum size in bytes of any single // compressed map response frame on the wire. If zero, // [DefaultMaxMessageSize] is used. MaxMessageSize int64}
// framedReader is an io.Reader that consumes a stream of length-prefixed// frames (each a little-endian uint32 length followed by that many bytes)// from r and yields only the frame payloads back-to-back.//// This lets us feed the concatenated zstd frames from our wire protocol// into a single streaming zstd decoder. Zstd's file format permits// concatenation (RFC 8478 §2), and klauspost's decoder handles it// transparently.//// If onNewFrame is non-nil, it is called after each new 4-byte length// header is successfully read. Used to reset the per-message decoded-size// budget downstream.type framedReader struct { r io.Reader maxSize int64 // per-frame compressed-size cap remain int // bytes remaining in the current frame onNewFrame func()}
func (f *framedReader) Read(p []byte) (int, error) { if f.remain == 0 { var hdr [4]byte if _, err := io.ReadFull(f.r, hdr[:]); err != nil { return 0, err } sz := int64(binary.LittleEndian.Uint32(hdr[:])) if sz == 0 { return 0, fmt.Errorf("map response: zero-length frame") } if sz > f.maxSize { return 0, fmt.Errorf("map response frame size %d exceeds max %d", sz, f.maxSize) } f.remain = int(sz) if f.onNewFrame != nil { f.onNewFrame() } } if len(p) > f.remain { p = p[:f.remain] } n, err := f.r.Read(p) f.remain -= n return n, err}
// boundedReader is an io.Reader that yields at most remain bytes from r// before returning an error. Call reset to raise the budget back to max,// typically at a new message boundary.//// Used to cap the decoded size of a single map response so a malicious// server can't send a small zstd frame that explodes into gigabytes of// junk for the json.Decoder to consume.type boundedReader struct { r io.Reader max int64 remain int64}
func (b *boundedReader) Read(p []byte) (int, error) { if b.remain <= 0 { return 0, fmt.Errorf("map response decoded size exceeds max %d", b.max) } if int64(len(p)) > b.remain { p = p[:b.remain] } n, err := b.r.Read(p) b.remain -= int64(n) return n, err}
func (b *boundedReader) reset() { b.remain = b.max }
// MapSession wraps an in-progress map response stream. Call Next to read// each MapResponse. Call Close when done.type MapSession struct { res *http.Response stream bool noiseDoer func(*http.Request) (*http.Response, error)
// inNext detects concurrent NextInto callers. It CAS-flips // false→true on entry and back to false on exit; a failed CAS // panics, akin to how the Go runtime detects concurrent map // access. It does not serialize Close vs. NextInto; that's // nextMu's job. inNext atomic.Bool
// nextMu is held while [MapSession.NextInto] is running jdec.Decode, // so that Close can wait for an in-flight Decode to unwind before it // touches zdec (Reset, pool-Put) and avoids racing with the running // Read chain that Decode drives. nextMu sync.Mutex read int // guarded by nextMu closed bool // guarded by nextMu zdec *zstd.Decoder // reads from a framedReader wrapping res.Body jdec *json.Decoder // reads decompressed JSON from zdec
closeOnce sync.Once closeErr error}
// NoiseRoundTrip sends an HTTP request over the Noise channel used by this map session.func (s *MapSession) NoiseRoundTrip(req *http.Request) (*http.Response, error) { return s.noiseDoer(req)}
// Next reads and returns the next MapResponse from the stream.// For non-streaming sessions, the first call returns the single response// and subsequent calls return io.EOF.// For streaming sessions, Next blocks until the next response arrives// or the server closes the connection.//// Each call allocates a fresh MapResponse. Callers that want to amortize// the allocation across calls can use [MapSession.NextInto].//// Next and NextInto are not safe to call concurrently from multiple// goroutines on the same [MapSession]; a concurrent call panics, akin// to the Go runtime's concurrent map access detection. [MapSession.Close]// may be called concurrently to abort an in-flight Next.func (s *MapSession) Next() (*tailcfg.MapResponse, error) { var resp tailcfg.MapResponse if err := s.NextInto(&resp); err != nil { return nil, err } return &resp, nil}
// NextInto is like [MapSession.Next] but decodes the next MapResponse into// the caller-supplied *resp rather than allocating a new one. The pointer's// pointee is zeroed before decoding so fields from a prior response do not// persist.//// For non-streaming sessions, the first call decodes the single response// and subsequent calls return io.EOF.// For streaming sessions, NextInto blocks until the next response arrives// or the server closes the connection.//// See [MapSession.Next] for concurrency rules; those apply to NextInto too.func (s *MapSession) NextInto(resp *tailcfg.MapResponse) error { if !s.inNext.CompareAndSwap(false, true) { panic("tsp: invalid concurrent call to MapSession.Next/NextInto") } defer s.inNext.Store(false)
s.nextMu.Lock() defer s.nextMu.Unlock() if s.closed { return errSessionClosed } if !s.stream && s.read > 0 { return io.EOF } *resp = tailcfg.MapResponse{} if err := s.jdec.Decode(resp); err != nil { return err } s.read++ return nil}
// Close returns the session's zstd decoder to the pool and closes the// underlying HTTP response body. It is safe to call Close multiple times// and from multiple goroutines, including while a [MapSession.Next] or// [MapSession.NextInto] call is in flight on another goroutine (which// will return an error once the body close propagates).func (s *MapSession) Close() error { // Callers are likely to race a deferred Close with a time.AfterFunc // timeout (or similar) Close that aborts a hung Next. Without the // Once, both Closes would Put the same *zstd.Decoder into the pool, // corrupting it, and the Reset/Put in one would race with the // zdec.Read that the hung Next is driving. // // Ordering inside the Once: close the body first to unblock any // in-flight NextInto (its Read chain ends at res.Body and will // return an error once it's closed). That lets NextInto unwind and // release nextMu. Only then do we take nextMu ourselves and touch // zdec, which is safe because no goroutine is still reading from // it. Acquiring nextMu before closing the body would deadlock // against a hung NextInto. s.closeOnce.Do(func() { s.closeErr = s.res.Body.Close() s.nextMu.Lock() defer s.nextMu.Unlock() s.closed = true s.zdec.Reset(nil) zstdDecoderPool.Put(s.zdec) }) return s.closeErr}
// SendMapUpdateOpts contains options for [Client.SendMapUpdate].type SendMapUpdateOpts struct { // NodeKey is the node's private key. Required. NodeKey key.NodePrivate
// DiscoKey, if non-zero, is the node's disco public key. // Peers use it to verify disco pings from this node, which is // what enables direct (non-DERP) paths. DiscoKey key.DiscoPublic
// Hostinfo is the host information to send. Optional; // if nil, a minimal default is used. Hostinfo *tailcfg.Hostinfo}
// SendMapUpdate sends a one-shot, non-streaming MapRequest to push small// updates (such as the node's endpoints, hostinfo, or disco public key) to the// coordination server without starting or disturbing a streaming map session.func (c *Client) SendMapUpdate(ctx context.Context, opts SendMapUpdateOpts) error { if opts.NodeKey.IsZero() { return fmt.Errorf("NodeKey is required") }
hi := opts.Hostinfo if hi == nil { hi = defaultHostinfo() }
mapReq := tailcfg.MapRequest{ Version: tailcfg.CurrentCapabilityVersion, NodeKey: opts.NodeKey.Public(), DiscoKey: opts.DiscoKey, Hostinfo: hi, Compress: "zstd",
// A lite update that lets the server persist our state without breaking // any existing streaming map session. See the [tailcfg.MapResponse] // OmitPeers docs. OmitPeers: true, Stream: false, ReadOnly: false, }
body, err := json.Marshal(mapReq) if err != nil { return fmt.Errorf("encoding map request: %w", err) }
nc, err := c.noiseClient(ctx) if err != nil { return fmt.Errorf("establishing noise connection: %w", err) }
url := c.serverURL + "/machine/map" url = strings.Replace(url, "http:", "https:", 1) req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body)) if err != nil { return fmt.Errorf("creating map request: %w", err) } ts2021.AddLBHeader(req, opts.NodeKey.Public())
res, err := nc.Do(req) if err != nil { return fmt.Errorf("map request: %w", err) } defer res.Body.Close()
if res.StatusCode != 200 { msg, _ := io.ReadAll(res.Body) return fmt.Errorf("map request: http %d: %.200s", res.StatusCode, strings.TrimSpace(string(msg))) } io.Copy(io.Discard, res.Body) return nil}
// Map sends a map request to the coordination server and returns a MapSession// for reading the framed, zstd-compressed response(s).func (c *Client) Map(ctx context.Context, opts MapOpts) (*MapSession, error) { if opts.NodeKey.IsZero() { return nil, fmt.Errorf("NodeKey is required") }
hi := opts.Hostinfo if hi == nil { hi = defaultHostinfo() }
mapReq := tailcfg.MapRequest{ Version: tailcfg.CurrentCapabilityVersion, NodeKey: opts.NodeKey.Public(), Hostinfo: hi, Stream: opts.Stream, Compress: "zstd", OmitPeers: opts.OmitPeers, // Streaming requires the server to track us as "connected", // which in turn requires ReadOnly=false. Non-streaming polls // stay ReadOnly to minimize side effects. ReadOnly: !opts.Stream, }
body, err := json.Marshal(mapReq) if err != nil { return nil, fmt.Errorf("encoding map request: %w", err) }
nc, err := c.noiseClient(ctx) if err != nil { return nil, fmt.Errorf("establishing noise connection: %w", err) }
url := c.serverURL + "/machine/map" url = strings.Replace(url, "http:", "https:", 1) // The response body is a stream of framed messages, each capped by // maxMessageSize below, but a streaming session can carry an unbounded // number of them over one body, so it must not be subject to // ts2021.Client.Do's default response size cap. ctx = httpbody.WithMaxSize(ctx, 0) req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("creating map request: %w", err) } ts2021.AddLBHeader(req, opts.NodeKey.Public())
res, err := nc.Do(req) if err != nil { return nil, fmt.Errorf("map request: %w", err) }
if res.StatusCode != 200 { // The body is an error message rather than a message stream, so // cap it: the map body itself is uncapped per the comment above. httpbody.LimitSizeTo(res, httpbody.DefaultMaxSize) msg, _ := io.ReadAll(res.Body) res.Body.Close() return nil, fmt.Errorf("map request: http %d: %.200s", res.StatusCode, strings.TrimSpace(string(msg))) }
maxMessageSize := cmp.Or(opts.MaxMessageSize, DefaultMaxMessageSize) bounded := &boundedReader{max: maxMessageSize, remain: maxMessageSize} fr := &framedReader{ r: res.Body, maxSize: maxMessageSize, onNewFrame: bounded.reset, }
zdec, _ := zstdDecoderPool.Get().(*zstd.Decoder) if zdec != nil { if err := zdec.Reset(fr); err != nil { // Reset can fail if the previous stream is in a bad state; drop // the decoder and create a fresh one. zdec = nil } } if zdec == nil { zdec, err = zstd.NewReader(fr, zstd.WithDecoderConcurrency(1)) if err != nil { res.Body.Close() return nil, fmt.Errorf("creating zstd decoder: %w", err) } } bounded.r = zdec
return &MapSession{ res: res, stream: opts.Stream, noiseDoer: nc.Do, zdec: zdec, jdec: json.NewDecoder(bounded), }, nil}