vendor: models/msg

models/msg/ctl.go

@@ -0,0 +1,46 @@
+// Copyright 2018 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package msg
+
+import (
+	"io"
+
+	jsonMsg "github.com/fatedier/golib/msg/json"
+)
+
+type Message = jsonMsg.Message
+
+var (
+	msgCtl *jsonMsg.MsgCtl
+)
+
+func init() {
+	msgCtl = jsonMsg.NewMsgCtl()
+	for typeByte, msg := range msgTypeMap {
+		msgCtl.RegisterMsg(typeByte, msg)
+	}
+}
+
+func ReadMsg(c io.Reader) (msg Message, err error) {
+	return msgCtl.ReadMsg(c)
+}
+
+func ReadMsgInto(c io.Reader, msg Message) (err error) {
+	return msgCtl.ReadMsgInto(c, msg)
+}
+
+func WriteMsg(c io.Writer, msg interface{}) (err error) {
+	return msgCtl.WriteMsg(c, msg)
+}

models/msg/msg.go

@@ -14,10 +14,7 @@
 
 package msg
 
-import (
-	"net"
-	"reflect"
-)
+import "net"
 
 const (
 	TypeLogin              = 'o'
@@ -40,39 +37,26 @@ const (
 )
 
 var (
-	TypeMap       map[byte]reflect.Type
-	TypeStringMap map[reflect.Type]byte
-)
-
-func init() {
-	TypeMap = make(map[byte]reflect.Type)
-	TypeStringMap = make(map[reflect.Type]byte)
-
-	TypeMap[TypeLogin] = reflect.TypeOf(Login{})
-	TypeMap[TypeLoginResp] = reflect.TypeOf(LoginResp{})
-	TypeMap[TypeNewProxy] = reflect.TypeOf(NewProxy{})
-	TypeMap[TypeNewProxyResp] = reflect.TypeOf(NewProxyResp{})
-	TypeMap[TypeCloseProxy] = reflect.TypeOf(CloseProxy{})
-	TypeMap[TypeNewWorkConn] = reflect.TypeOf(NewWorkConn{})
-	TypeMap[TypeReqWorkConn] = reflect.TypeOf(ReqWorkConn{})
-	TypeMap[TypeStartWorkConn] = reflect.TypeOf(StartWorkConn{})
-	TypeMap[TypeNewVisitorConn] = reflect.TypeOf(NewVisitorConn{})
-	TypeMap[TypeNewVisitorConnResp] = reflect.TypeOf(NewVisitorConnResp{})
-	TypeMap[TypePing] = reflect.TypeOf(Ping{})
-	TypeMap[TypePong] = reflect.TypeOf(Pong{})
-	TypeMap[TypeUdpPacket] = reflect.TypeOf(UdpPacket{})
-	TypeMap[TypeNatHoleVisitor] = reflect.TypeOf(NatHoleVisitor{})
-	TypeMap[TypeNatHoleClient] = reflect.TypeOf(NatHoleClient{})
-	TypeMap[TypeNatHoleResp] = reflect.TypeOf(NatHoleResp{})
-	TypeMap[TypeNatHoleSid] = reflect.TypeOf(NatHoleSid{})
-
-	for k, v := range TypeMap {
-		TypeStringMap[v] = k
+	msgTypeMap = map[byte]interface{}{
+		TypeLogin:              Login{},
+		TypeLoginResp:          LoginResp{},
+		TypeNewProxy:           NewProxy{},
+		TypeNewProxyResp:       NewProxyResp{},
+		TypeCloseProxy:         CloseProxy{},
+		TypeNewWorkConn:        NewWorkConn{},
+		TypeReqWorkConn:        ReqWorkConn{},
+		TypeStartWorkConn:      StartWorkConn{},
+		TypeNewVisitorConn:     NewVisitorConn{},
+		TypeNewVisitorConnResp: NewVisitorConnResp{},
+		TypePing:               Ping{},
+		TypePong:               Pong{},
+		TypeUdpPacket:          UdpPacket{},
+		TypeNatHoleVisitor:     NatHoleVisitor{},
+		TypeNatHoleClient:      NatHoleClient{},
+		TypeNatHoleResp:        NatHoleResp{},
+		TypeNatHoleSid:         NatHoleSid{},
 	}
-}
-
-// Message wraps socket packages for communicating between frpc and frps.
-type Message interface{}
+)
 
 // When frpc start, client send this message to login to server.
 type Login struct {

models/msg/pack.go

@@ -1,69 +0,0 @@
-// Copyright 2016 fatedier, fatedier@gmail.com
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package msg
-
-import (
-	"bytes"
-	"encoding/binary"
-	"encoding/json"
-	"fmt"
-	"reflect"
-
-	"github.com/fatedier/frp/models/errors"
-)
-
-func unpack(typeByte byte, buffer []byte, msgIn Message) (msg Message, err error) {
-	if msgIn == nil {
-		t, ok := TypeMap[typeByte]
-		if !ok {
-			err = fmt.Errorf("Unsupported message type %b", typeByte)
-			return
-		}
-
-		msg = reflect.New(t).Interface().(Message)
-	} else {
-		msg = msgIn
-	}
-
-	err = json.Unmarshal(buffer, &msg)
-	return
-}
-
-func UnPackInto(buffer []byte, msg Message) (err error) {
-	_, err = unpack(' ', buffer, msg)
-	return
-}
-
-func UnPack(typeByte byte, buffer []byte) (msg Message, err error) {
-	return unpack(typeByte, buffer, nil)
-}
-
-func Pack(msg Message) ([]byte, error) {
-	typeByte, ok := TypeStringMap[reflect.TypeOf(msg).Elem()]
-	if !ok {
-		return nil, errors.ErrMsgType
-	}
-
-	content, err := json.Marshal(msg)
-	if err != nil {
-		return nil, err
-	}
-
-	buffer := bytes.NewBuffer(nil)
-	buffer.WriteByte(typeByte)
-	binary.Write(buffer, binary.BigEndian, int64(len(content)))
-	buffer.Write(content)
-	return buffer.Bytes(), nil
-}

models/msg/pack_test.go

@@ -1,87 +0,0 @@
-// Copyright 2016 fatedier, fatedier@gmail.com
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package msg
-
-import (
-	"bytes"
-	"encoding/binary"
-	"reflect"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-
-	"github.com/fatedier/frp/models/errors"
-)
-
-type TestStruct struct{}
-
-func TestPack(t *testing.T) {
-	assert := assert.New(t)
-
-	var (
-		msg    Message
-		buffer []byte
-		err    error
-	)
-
-	// error type
-	msg = &TestStruct{}
-	buffer, err = Pack(msg)
-	assert.Error(err, errors.ErrMsgType.Error())
-
-	// correct
-	msg = &Ping{}
-	buffer, err = Pack(msg)
-	assert.NoError(err)
-	b := bytes.NewBuffer(nil)
-	b.WriteByte(TypePing)
-	binary.Write(b, binary.BigEndian, int64(2))
-	b.WriteString("{}")
-	assert.True(bytes.Equal(b.Bytes(), buffer))
-}
-
-func TestUnPack(t *testing.T) {
-	assert := assert.New(t)
-
-	var (
-		msg Message
-		err error
-	)
-
-	// error message type
-	msg, err = UnPack('-', []byte("{}"))
-	assert.Error(err)
-
-	// correct
-	msg, err = UnPack(TypePong, []byte("{}"))
-	assert.NoError(err)
-	assert.Equal(reflect.TypeOf(msg).Elem(), reflect.TypeOf(Pong{}))
-}
-
-func TestUnPackInto(t *testing.T) {
-	assert := assert.New(t)
-
-	var err error
-
-	// correct type
-	pongMsg := &Pong{}
-	err = UnPackInto([]byte("{}"), pongMsg)
-	assert.NoError(err)
-
-	// wrong type
-	loginMsg := &Login{}
-	err = UnPackInto([]byte(`{"version": 123}`), loginMsg)
-	assert.Error(err)
-}

models/msg/process.go

@@ -1,88 +0,0 @@
-// Copyright 2016 fatedier, fatedier@gmail.com
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package msg
-
-import (
-	"encoding/binary"
-	"fmt"
-	"io"
-)
-
-var (
-	MaxMsgLength int64 = 10240
-)
-
-func readMsg(c io.Reader) (typeByte byte, buffer []byte, err error) {
-	buffer = make([]byte, 1)
-	_, err = c.Read(buffer)
-	if err != nil {
-		return
-	}
-	typeByte = buffer[0]
-	if _, ok := TypeMap[typeByte]; !ok {
-		err = fmt.Errorf("Message type error")
-		return
-	}
-
-	var length int64
-	err = binary.Read(c, binary.BigEndian, &length)
-	if err != nil {
-		return
-	}
-	if length > MaxMsgLength {
-		err = fmt.Errorf("Message length exceed the limit")
-		return
-	}
-
-	buffer = make([]byte, length)
-	n, err := io.ReadFull(c, buffer)
-	if err != nil {
-		return
-	}
-
-	if int64(n) != length {
-		err = fmt.Errorf("Message format error")
-	}
-	return
-}
-
-func ReadMsg(c io.Reader) (msg Message, err error) {
-	typeByte, buffer, err := readMsg(c)
-	if err != nil {
-		return
-	}
-	return UnPack(typeByte, buffer)
-}
-
-func ReadMsgInto(c io.Reader, msg Message) (err error) {
-	_, buffer, err := readMsg(c)
-	if err != nil {
-		return
-	}
-	return UnPackInto(buffer, msg)
-}
-
-func WriteMsg(c io.Writer, msg interface{}) (err error) {
-	buffer, err := Pack(msg)
-	if err != nil {
-		return
-	}
-
-	if _, err = c.Write(buffer); err != nil {
-		return
-	}
-
-	return nil
-}

models/msg/process_test.go

@@ -1,97 +0,0 @@
-// Copyright 2016 fatedier, fatedier@gmail.com
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package msg
-
-import (
-	"bytes"
-	"encoding/binary"
-	"reflect"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-)
-
-func TestProcess(t *testing.T) {
-	assert := assert.New(t)
-
-	var (
-		msg    Message
-		resMsg Message
-		err    error
-	)
-	// empty struct
-	msg = &Ping{}
-	buffer := bytes.NewBuffer(nil)
-	err = WriteMsg(buffer, msg)
-	assert.NoError(err)
-
-	resMsg, err = ReadMsg(buffer)
-	assert.NoError(err)
-	assert.Equal(reflect.TypeOf(resMsg).Elem(), TypeMap[TypePing])
-
-	// normal message
-	msg = &StartWorkConn{
-		ProxyName: "test",
-	}
-	buffer = bytes.NewBuffer(nil)
-	err = WriteMsg(buffer, msg)
-	assert.NoError(err)
-
-	resMsg, err = ReadMsg(buffer)
-	assert.NoError(err)
-	assert.Equal(reflect.TypeOf(resMsg).Elem(), TypeMap[TypeStartWorkConn])
-
-	startWorkConnMsg, ok := resMsg.(*StartWorkConn)
-	assert.True(ok)
-	assert.Equal("test", startWorkConnMsg.ProxyName)
-
-	// ReadMsgInto correct
-	msg = &Pong{}
-	buffer = bytes.NewBuffer(nil)
-	err = WriteMsg(buffer, msg)
-	assert.NoError(err)
-
-	err = ReadMsgInto(buffer, msg)
-	assert.NoError(err)
-
-	// ReadMsgInto error type
-	content := []byte(`{"run_id": 123}`)
-	buffer = bytes.NewBuffer(nil)
-	buffer.WriteByte(TypeNewWorkConn)
-	binary.Write(buffer, binary.BigEndian, int64(len(content)))
-	buffer.Write(content)
-
-	resMsg = &NewWorkConn{}
-	err = ReadMsgInto(buffer, resMsg)
-	assert.Error(err)
-
-	// message format error
-	buffer = bytes.NewBuffer([]byte("1234"))
-
-	resMsg = &NewProxyResp{}
-	err = ReadMsgInto(buffer, resMsg)
-	assert.Error(err)
-
-	// MaxLength, real message length is 2
-	MaxMsgLength = 1
-	msg = &Ping{}
-	buffer = bytes.NewBuffer(nil)
-	err = WriteMsg(buffer, msg)
-	assert.NoError(err)
-
-	_, err = ReadMsg(buffer)
-	assert.Error(err)
-	return
-}