go-tunnel/control/message.go

74 lines
1.4 KiB
Go

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
}