chore: initial commit

control/auth.go

@@ -0,0 +1,9 @@
+package control
+
+type AuthRequestPayload struct {
+	Credentials interface{} `json:"c"`
+}
+
+type AuthResponsePayload struct {
+	Success bool `json:"s"`
+}

control/control.go

@@ -0,0 +1,167 @@
+package control
+
+import (
+	"context"
+	"encoding/json"
+	"net"
+	"time"
+
+	"github.com/pkg/errors"
+	"github.com/xtaci/smux"
+	"gitlab.com/wpetit/goweb/logger"
+)
+
+type Control struct {
+	encoder *json.Encoder
+	decoder *json.Decoder
+	stream  *smux.Stream
+	sess    *smux.Session
+}
+
+func (c *Control) AuthRequest(credentials interface{}) (bool, error) {
+	req := NewMessage(TypeAuthRequest, &AuthRequestPayload{
+		Credentials: credentials,
+	})
+
+	res := NewMessage(TypeAuthResponse, nil)
+
+	if err := c.reqRes(req, res); err != nil {
+		return false, errors.WithStack(err)
+	}
+
+	authResPayload, ok := res.Payload.(*AuthResponsePayload)
+	if !ok {
+		return false, errors.WithStack(ErrUnexpectedMessage)
+	}
+
+	return authResPayload.Success, nil
+}
+
+type CloseStream func()
+
+func (c *Control) Proxy(ctx context.Context, network, address string) (net.Conn, error) {
+	var (
+		stream *smux.Stream
+		err    error
+	)
+
+	req := NewMessage(TypeProxyRequest, &ProxyRequestPayload{
+		Network: network,
+		Address: address,
+	})
+
+	ctx = logger.With(ctx, logger.F("network", network), logger.F("address", address))
+
+	logger.Debug(ctx, "proxying")
+
+	if err := c.Write(req); err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	logger.Debug(ctx, "opening stream")
+
+	stream, err = c.sess.AcceptStream()
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	return stream, nil
+}
+
+func (c *Control) Listen(ctx context.Context, handlers Handlers) error {
+	for {
+		logger.Debug(ctx, "reading next message")
+
+		req, err := c.Read()
+		if err != nil {
+			return errors.WithStack(err)
+		}
+
+		go func() {
+			subCtx := logger.With(ctx, logger.F("messageType", req.Type))
+
+			handler, exists := handlers[req.Type]
+			if !exists {
+				logger.Error(subCtx, "no message handler registered")
+
+				return
+			}
+
+			res, err := handler(subCtx, req)
+			if err != nil {
+				logger.Error(subCtx, "error while handling message", logger.E(err))
+
+				return
+			}
+
+			if res == nil {
+				return
+			}
+
+			if err := c.Write(res); err != nil {
+				logger.Error(subCtx, "error while write message response", logger.E(err))
+
+				return
+			}
+		}()
+	}
+}
+
+func (c *Control) Read() (*Message, error) {
+	message := &Message{}
+
+	if err := c.read(message); err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	return message, nil
+}
+
+func (c *Control) Write(m *Message) error {
+	if err := c.write(m); err != nil {
+		return errors.WithStack(err)
+	}
+
+	return nil
+}
+
+func (c *Control) reqRes(req *Message, res *Message) error {
+	if err := c.write(req); err != nil {
+		return errors.WithStack(err)
+	}
+
+	if err := c.read(res); err != nil {
+		return errors.WithStack(err)
+	}
+
+	return nil
+}
+
+func (c *Control) read(m *Message) error {
+	if err := c.decoder.Decode(m); err != nil {
+		return errors.WithStack(err)
+	}
+
+	return nil
+}
+
+func (c *Control) write(m *Message) error {
+	if err := c.stream.SetWriteDeadline(time.Now().Add(time.Second)); err != nil {
+		return errors.WithStack(err)
+	}
+
+	if err := c.encoder.Encode(m); err != nil {
+		return errors.WithStack(err)
+	}
+
+	return nil
+}
+
+func New(sess *smux.Session, controlStream *smux.Stream) *Control {
+	return &Control{
+		encoder: json.NewEncoder(controlStream),
+		decoder: json.NewDecoder(controlStream),
+		sess:    sess,
+		stream:  controlStream,
+	}
+}

control/error.go

@@ -0,0 +1,7 @@
+package control
+
+import "errors"
+
+var (
+	ErrUnexpectedMessage = errors.New("unexpected message")
+)

control/handler.go

@@ -0,0 +1,7 @@
+package control
+
+import "context"
+
+type Handlers map[MessageType]MessageHandler
+
+type MessageHandler func(ctx context.Context, m *Message) (*Message, error)

control/message.go

@@ -0,0 +1,73 @@
+package control
+
+import (
+	"encoding/json"
+
+	"github.com/pkg/errors"
+)
+
+const (
+	TypeAuthRequest  MessageType = "auth-req"
+	TypeAuthResponse MessageType = "auth-res"
+	TypeProxyRequest MessageType = "proxy-req"
+)
+
+type MessageType string
+
+type BaseMessage struct {
+	Type       MessageType     `json:"t"`
+	RawPayload json.RawMessage `json:"p"`
+}
+
+type Message struct {
+	BaseMessage
+	Payload interface{} `json:"p"`
+}
+
+func (m *Message) UnmarshalJSON(data []byte) error {
+	base := &BaseMessage{}
+
+	if err := json.Unmarshal(data, base); err != nil {
+		return errors.WithStack(err)
+	}
+
+	payload, err := unmarshalPayload(base.Type, base.RawPayload)
+	if err != nil {
+		return errors.WithStack(err)
+	}
+
+	m.Type = base.Type
+	m.Payload = payload
+
+	return nil
+}
+
+func NewMessage(mType MessageType, payload interface{}) *Message {
+	return &Message{
+		BaseMessage: BaseMessage{
+			Type: mType,
+		},
+		Payload: payload,
+	}
+}
+
+func unmarshalPayload(mType MessageType, data []byte) (interface{}, error) {
+	var payload interface{}
+
+	switch mType {
+	case TypeAuthRequest:
+		payload = &AuthRequestPayload{}
+	case TypeAuthResponse:
+		payload = &AuthResponsePayload{}
+	case TypeProxyRequest:
+		payload = &ProxyRequestPayload{}
+	default:
+		return nil, errors.Wrapf(ErrUnexpectedMessage, "unexpected message type '%s'", mType)
+	}
+
+	if err := json.Unmarshal(data, payload); err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	return payload, nil
+}

control/proxy.go

@@ -0,0 +1,6 @@
+package control
+
+type ProxyRequestPayload struct {
+	Network string `json:"n"`
+	Address string `json:"a"`
+}