diff --git a/models/msg/ctl.go b/models/msg/ctl.go new file mode 100644 index 0000000..0eafdbc --- /dev/null +++ b/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) +} diff --git a/models/msg/msg.go b/models/msg/msg.go index 0cdb3f4..93214fc 100644 --- a/models/msg/msg.go +++ b/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 { diff --git a/models/msg/pack.go b/models/msg/pack.go deleted file mode 100644 index afb34ae..0000000 --- a/models/msg/pack.go +++ /dev/null @@ -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 -} diff --git a/models/msg/pack_test.go b/models/msg/pack_test.go deleted file mode 100644 index 6aab682..0000000 --- a/models/msg/pack_test.go +++ /dev/null @@ -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) -} diff --git a/models/msg/process.go b/models/msg/process.go deleted file mode 100644 index d5a89d9..0000000 --- a/models/msg/process.go +++ /dev/null @@ -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 -} diff --git a/models/msg/process_test.go b/models/msg/process_test.go deleted file mode 100644 index 149b7db..0000000 --- a/models/msg/process_test.go +++ /dev/null @@ -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 -}