wstunnel

README

wstunnel

Shared WebSocket + yamux tunnel used across the Miabi ecosystem: one outbound, NAT-friendly WebSocket carries many independent multiplexed streams. It is the single source of truth for the wire protocol so the control plane and the agents/runners that dial into it always agree on framing and keepalive — the two ends drifting is what silently breaks a tunnel, which is exactly what this module prevents.

It depends only on gorilla/websocket and hashicorp/yamux (no logging, ORM, or Docker deps), so the lean agent/runner binaries stay small.

Roles

• Control plane opens streams → wstunnel.Client(ws)
• Agent / runner accepts streams → wstunnel.Server(ws) (or Dial/Serve)

Usage

Control plane (accepts the WebSocket, drives streams):

ws, _ := upgrader.Upgrade(w, r, nil)
sess, _ := wstunnel.Client(ws)
stream, _ := sess.OpenStream() // one per request / job lease

Agent or runner (dials out, services accepted streams, auto-reconnects):

opts := wstunnel.ClientOptions{
    URL:    wstunnel.URL(controlURL, "/api/v1/runner/connect"),
    Header: http.Header{"Authorization": {"Bearer " + token}},
    OnConnect: func() { log.Info("connected") },
    OnError:   func(err error) { log.Warn("disconnected", err) },
}
_ = wstunnel.Serve(ctx, opts, func(ctx context.Context, sess *yamux.Session) error {
    for {
        stream, err := sess.AcceptStream()
        if err != nil {
            return err
        }
        go handle(stream) // pipe to Docker, service a job lease, etc.
    }
})

API

Symbol	Purpose
NewConn(ws) net.Conn	Adapt a *websocket.Conn to a net.Conn carrying yamux framing
Config() *yamux.Config	The shared keepalive/timeout config both ends use
Client(ws) / Server(ws)	Open the stream-opening / stream-accepting session
URL(base, path)	Normalize an http(s) base + path to the ws(s) endpoint
Dial(ctx, opts)	One-shot dial → accepting session
Serve(ctx, opts, handle)	Reconnecting accept loop with exponential backoff

Licensed under Apache-2.0.

Documentation

Overview

Package wstunnel adapts a single WebSocket into a multiplexed byte stream (yamux over WebSocket), so one outbound, NAT-friendly connection can carry many independent streams. It is the shared wire protocol for the Miabi control plane and the agents/runners that dial into it: the control plane OPENS streams (Client), the agent/runner ACCEPTS them (Server). Both ends MUST use this package so the yamux framing and keepalive settings match exactly — a mismatch silently breaks the tunnel.

Index

• Constants
• func Client(ws *websocket.Conn) (*yamux.Session, error)
• func Config() *yamux.Config
• func Dial(ctx context.Context, opts ClientOptions) (*yamux.Session, error)
• func NewConn(ws *websocket.Conn) net.Conn
• func Serve(ctx context.Context, opts ClientOptions, ...) error
• func Server(ws *websocket.Conn) (*yamux.Session, error)
• func URL(base, path string) string
• type ClientOptions

Constants

const (
	KeepAliveInterval      = 20 * time.Second
	ConnectionWriteTimeout = 15 * time.Second
)

Keepalive/timeout defaults for the yamux session. Exported so callers can build a custom Config from the same baseline if needed.

Functions

func Client

func Client(ws *websocket.Conn) (*yamux.Session, error)

Client opens the stream-opening side of the session over ws (the control plane): it OPENS streams (e.g. one per Docker request or job lease).

func Config

func Config() *yamux.Config

Config returns the shared yamux config: keepalive enabled and logging muted. Both ends use the same values, so the framing and liveness detection agree.

func Dial

func Dial(ctx context.Context, opts ClientOptions) (*yamux.Session, error)

Dial opens one WebSocket to opts.URL and returns the accepting (Server-side) yamux session. The caller owns closing the session (which closes the ws).

func NewConn

func NewConn(ws *websocket.Conn) net.Conn

NewConn wraps a WebSocket connection as a net.Conn carrying the yamux framing.

func Serve

func Serve(ctx context.Context, opts ClientOptions, handle func(ctx context.Context, sess *yamux.Session) error) error

Serve runs the accepting side until ctx is cancelled, reconnecting with exponential backoff. For each live session it calls handle(ctx, sess), which should block for the session's lifetime (e.g. accepting and servicing streams) and return when it ends. A clean session (handle returns nil) resets the backoff. handle owns the session; Serve closes it before retrying.

func Server

func Server(ws *websocket.Conn) (*yamux.Session, error)

Server opens the stream-accepting side of the session over ws (an agent or runner): it ACCEPTS streams opened by the control plane.

func URL

func URL(base, path string) string

URL converts an http(s) control-plane base URL and a connect path into the ws(s) endpoint a client dials. It normalizes the scheme (http→ws, https→wss) and joins the path with exactly one slash, so callers pass their own connect path (e.g. "/api/v1/agent/connect" or "/api/v1/runner/connect").

Types

type ClientOptions

type ClientOptions struct {
	// URL is the fully-resolved ws(s) endpoint (see URL()).
	URL string
	// Header carries auth and metadata on the WebSocket handshake (e.g.
	// "Authorization: Bearer <token>", "X-Runner-Version").
	Header http.Header
	// Dialer overrides the WebSocket dialer (nil uses a copy of the default).
	Dialer *websocket.Dialer
	// Insecure skips TLS verification (dev only). Ignored when Dialer is set.
	Insecure bool
	// MinBackoff / MaxBackoff bound the reconnect delay (defaults 1s / 30s).
	MinBackoff, MaxBackoff time.Duration
	// OnConnect fires after each successful session is established; OnError fires
	// with the cause each time a session ends or a dial fails. Both are optional
	// (nil = no-op), keeping this package free of any logging dependency.
	OnConnect func()
	OnError   func(error)
}

ClientOptions configures the accepting (agent/runner) side of the tunnel.