route/lib/tun2/connection.go

package tun2

import (
	"bufio"
	"context"
	"net"
	"net/http"
	"time"

	"github.com/Xe/ln"
	failure "github.com/dgryski/go-failure"
	"github.com/pkg/errors"
	"github.com/xtaci/smux"
)

// Connection is a single active client -> server connection and session
// containing many streams over TCP+TLS or KCP+TLS. Every stream beyond the
// control stream is assumed to be passed to the underlying backend server.
type Connection struct {
	id            string
	conn          net.Conn
	isKCP         bool
	session       *smux.Session
	controlStream *smux.Stream
	user          string
	domain        string
	cancel        context.CancelFunc
	detector      *failure.Detector
	Auth          *Auth
}

// F logs key->value pairs as an ln.Fer
func (c *Connection) F() ln.F {
	return map[string]interface{}{
		"id":     c.id,
		"remote": c.conn.RemoteAddr(),
		"local":  c.conn.LocalAddr(),
		"isKCP":  c.isKCP,
		"user":   c.user,
		"domain": c.domain,
	}
}

// Ping ends a "ping" to the client. If the client doesn't respond or the connection
// dies, then the connection needs to be cleaned up.
func (c *Connection) Ping() error {
	req, err := http.NewRequest("GET", "http://backend/health", nil)
	if err != nil {
		panic(err)
	}

	_, err = c.RoundTrip(req)
	if err != nil {
		ln.Error(err, c.F(), ln.F{"action": "ping_roundtrip"})
		defer c.cancel()
		return err
	}

	c.detector.Ping(time.Now())

	return nil
}

// OpenStream creates a new stream (connection) to the backend server.
func (c *Connection) OpenStream() (net.Conn, error) {
	err := c.conn.SetDeadline(time.Now().Add(time.Second))
	if err != nil {
		ln.Error(err, c.F())
		return nil, err
	}

	stream, err := c.session.OpenStream()
	if err != nil {
		ln.Error(err, c.F())
		return nil, err
	}

	return stream, c.conn.SetDeadline(time.Time{})
}

// Close destroys resouces specific to the connection.
func (c *Connection) Close() error {
	err := c.controlStream.Close()
	if err != nil {
		return err
	}

	err = c.session.Close()
	if err != nil {
		return err
	}

	err = c.conn.Close()
	if err != nil {
		return err
	}

	return nil
}

// Connection-specific errors
var (
	ErrCantOpenSessionStream = errors.New("tun2: connection can't open session stream")
	ErrCantWriteRequest      = errors.New("tun2: connection stream can't write request")
	ErrCantReadResponse      = errors.New("tun2: connection stream can't read response")
)

// RoundTrip forwards a HTTP request to the remote backend and then returns the
// response, if any.
func (c *Connection) RoundTrip(req *http.Request) (*http.Response, error) {
	stream, err := c.OpenStream()
	if err != nil {
		return nil, errors.Wrap(err, ErrCantOpenSessionStream.Error())
	}
	defer stream.Close()

	err = req.Write(stream)
	if err != nil {
		return nil, errors.Wrap(err, ErrCantWriteRequest.Error())
	}

	buf := bufio.NewReader(stream)

	resp, err := http.ReadResponse(buf, req)
	if err != nil {
		return nil, errors.Wrap(err, ErrCantReadResponse.Error())
	}

	return resp, nil
}