pack.go 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. // Copyright 2016 fatedier, fatedier@gmail.com
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package msg
  15. import (
  16. "bytes"
  17. "encoding/binary"
  18. "encoding/json"
  19. "fmt"
  20. "reflect"
  21. "github.com/fatedier/frp/utils/errors"
  22. )
  23. func unpack(typeByte byte, buffer []byte, msgIn Message) (msg Message, err error) {
  24. if msgIn == nil {
  25. t, ok := TypeMap[typeByte]
  26. if !ok {
  27. err = fmt.Errorf("Unsupported message type %b", typeByte)
  28. return
  29. }
  30. msg = reflect.New(t).Interface().(Message)
  31. } else {
  32. msg = msgIn
  33. }
  34. err = json.Unmarshal(buffer, &msg)
  35. return
  36. }
  37. func UnPackInto(buffer []byte, msg Message) (err error) {
  38. _, err = unpack(' ', buffer, msg)
  39. return
  40. }
  41. func UnPack(typeByte byte, buffer []byte) (msg Message, err error) {
  42. return unpack(typeByte, buffer, nil)
  43. }
  44. func Pack(msg Message) ([]byte, error) {
  45. typeByte, ok := TypeStringMap[reflect.TypeOf(msg).Elem()]
  46. if !ok {
  47. return nil, errors.ErrMsgType
  48. }
  49. content, err := json.Marshal(msg)
  50. if err != nil {
  51. return nil, err
  52. }
  53. buffer := bytes.NewBuffer(nil)
  54. buffer.WriteByte(typeByte)
  55. binary.Write(buffer, binary.BigEndian, int64(len(content)))
  56. buffer.Write(content)
  57. return buffer.Bytes(), nil
  58. }