Goを使ったメッセージのエンコードとデコードのコード設計

メッセージをバイト列にエンコードし、バイト列からデコードするためのGoのコード設計を考えたい。メッセージにはたくさんの種類があり、それぞれペイロードが異なるとする。EncodeとDecodeの2つのAPIを設計することとする。

まず、メッセージのフォーマットはどう設計すべきか？最もシンプルなのはペイロードのサイズ、メッセージの種類、ペイロードの実体をバイト列に並べる構成だ。

[4-byte payload length][1-byte message type][N-byte payload]

このメッセージを、例えばTCPを使って受け取る場合、一般的にはバイト列のストリームを受け取って処理することになる。受け取ったバイト列には、一つのメッセージが含まれるかもしれないし、中途半端に受け取っているかもしれないし、複数のメッセージが埋め込まれているかもしれない。

// 送信したいメッセージ AAA, BBB

[AA] // 途中
[AAA] // AAAがすべて届いた
[AAABB] // Bが途中
[AAABBB]

したがって、メッセージフレームを読み取る・書き込む処理が必要そうだ。

func readFrame(r io.Reader) ([]byte, error) {}
func writeFrame(w io.Writer, messageType byte, payload []byte) error {}

readFrameには以下のステップの処理が必要になる。

4バイト読む。
4バイトを整数に変換し、Nとする。
1+Nバイトを読む。
フレームを組み立てて返す。

const MaxPayloadSize = 1 << 20 // 1 MiB

func readFrame(r io.Reader) ([]byte, error) {
  var header [4]byte

  if _, err := io.ReadFull(r, header[:]); err != nil {
    return nil, err
  }
  payloadSize := binary.BigEndian.Uint32(header[:])
  if payloadSize > MaxPayloadSize {
    return nil, fmt.Errorf("payload too large: %d", payloadSize)
  }
  frame := make([]byte, 4+1+int(payloadSize))
  copy(frame[:4], header[:])
  if _, err := io.ReadFull(r, frame[4:]); err != nil {
    return nil, err
  }
  return frame, nil
}

メモ:

io.ReadFullはスライスが埋まるまで内部でReadを繰り返す。
header[:]は配列であるheaderを参照するスライスに変換している。
無条件でmakeを呼ぶと不必要にメモリを確保することにつながるので、必ずその前に最大ペイロード長をチェックしている。

writeFrameには以下のステップが必要になる。

メッセージを検証する。ペイロードのサイズをチェックし、大きすぎる場合はエラーを返す。
フレームを組み立てる。
フレームをストリームに書き込む。

func writeFrame(w io.Writer, messageType byte, payload []byte) error {
  if len(payload) > MaxPayloadSize {
    return fmt.Errorf("payload too large: %d", len(payload))
  }
  frame := make([]byte, 4+1+len(payload))
  binary.BigEndian.PutUint32(frame[:4], uint32(len(payload)))
  frame[4] = messageType
  copy(frame[5:], payload)
  return writeAll(w, frame)
}

func writeAll(w io.Writer, data []byte) error {
  for len(data) > 0 {
    n, err := w.Write(data)
    if err != nil {
      return err
    }
    if n == 0 {
      return io.ErrShortWrite
    }
    data = data[n:]
  }
  return nil
}

メモ:

io.WriterのWriteは、渡されたデータより少ないバイト数しか書き込めなかった場合、非nilのエラーも返す必要がある。一方、Readは一部だけ読み取ってエラーを返さないことがある。そのため、標準ライブラリにはReadFullがあるが、WriteFullはない。
w.Writeは、dataの先頭から書き込んだバイト数とエラーを返す。
n == 0の場合は、無限ループに陥ることを避けるため、ループを続けずにエラーを返す。

Messageはどう設計できるだろうか。メッセージには複数の種類が存在し、それぞれ異なるエンコード、デコードの方法を持つと考えられる。

メッセージタイプ
ペイロードのエンコード

メッセージごとに保持するデータやペイロードのエンコード方法は異なる。一方、すべてのメッセージは、メッセージタイプを持ち、ペイロードへエンコードできるという共通の振る舞いを持つ。だから、これらの振る舞いをインターフェースとして定義するのが自然だ。

type MessageType uint8
type Message interface {
  Type() MessageType
  EncodePayload() ([]byte, error)
}

const (
  MessageTypeUnknown MessageType = iota
  MessageTypeText
)

type TextMessage struct {
  Text string
}

func (TextMessage) Type() MessageType {
  return MessageTypeText
}

func (m TextMessage) EncodePayload() ([]byte, error) {
  return []byte(m.Text), nil
}

func Encode(message Message) ([]byte, error) {
  if message == nil {
    return nil, errors.New("message is nil")
  }
  payload, err := message.EncodePayload()
  if err != nil {
    return nil, fmt.Errorf("encode message %d: %w", message.Type(), err)
  }
  if len(payload) > MaxPayloadSize {
    return nil, fmt.Errorf("payload too large: %d", len(payload))
  }
  return payload, nil
}

デコードはメッセージごとに異なるロジックとなる。また、Messageがデコードをするのではなく、バイト列を受け取って処理するため、メッセージの振る舞いとして定義するのは不自然だ。テキストメッセージのデコードロジックをまず定義してみよう。

type MessageDecoder func(payload []byte) (Message, error)

func decodeTextMessage(payload []byte) (Message, error) {
  message := &TextMessage{
    Text: string(payload),
  }
  return message, nil
}

メッセージタイプと対応するデコーダーは、ハッシュマップを使って管理できる。ファクトリーパターンを使うことで、呼び出し時にはメッセージタイプとペイロードがあればデコードできるようになる。

const frameHeaderSize = 5

var decoders = map[MessageType]MessageDecoder{
  MessageTypeText: decodeTextMessage,
}

func Decode(frame []byte) (Message, error) {
  // frame check
  if len(frame) < frameHeaderSize {
    return nil, fmt.Errorf("frame too short: %d", len(frame))
  }

  // validate payload size
  payloadSize := binary.BigEndian.Uint32(frame[:4])
  if payloadSize > MaxPayloadSize {
    return nil, fmt.Errorf("payload too large: %d", payloadSize)
  }
  actualPayloadSize := len(frame) - frameHeaderSize
  if actualPayloadSize != int(payloadSize) {
    return nil, fmt.Errorf("invalid payload size: expected %d, got %d", payloadSize, actualPayloadSize)
  }
  messageType := MessageType(frame[4])
  decoder, ok := decoders[messageType]
  if !ok {
    return nil, fmt.Errorf("unknown message type: %d", messageType)
  }

  message, err := decoder(frame[frameHeaderSize:])
  if err != nil {
    return nil, fmt.Errorf("decode message type %d: %w", messageType, err)
  }
  return message, nil
}