Goでこのパケットの読み書きを行う

このようなパケットの読み書き処理をGoで素直に書くと次のようになります。

// 1パケット分読み出す
func readPacket(r io.Reader) ([]byte, error) {
    var size uint16
    if err := binary.Read(r, binary.LittleEndian, &size); err != nil {
        return nil, err
    }
    p := make([]byte, size)
    if _, err := io.ReadFull(r, p); err != nil {
        return nil, err
    }
    return p, nil
}

// 1パケット分書き出す
func writePacket(w io.Writer, p []byte) error {
    if err := binary.Write(w, binary.LittleEndian, uint16(len(p))); err != nil {
        return err
    }
    if _, err := w.Write(p); err != nil {
        return err
    }
    return nil
}

これを使って、例えば "hello, world"という文字列をパケットにすると次のようになります。

func main() {
    var b bytes.Buffer
    payload := []byte("hello, world")
    writePacket(&b, payload)
    print(hex.Dump(b.Bytes()))
}

先頭に 0c 00（payloadの長さ）が付加されていますね！

00000000  0c 00 68 65 6c 6c 6f 2c  20 77 6f 72 6c 64        |..hello, world|

パケット構造の仕様が増える場合

ここで、「Payloadを暗号化してほしい！」という追加の仕様がきたとします。
そうなると暗号化の処理が追加で必要となります。

暗号化の関数を encrypt() とすると、パケットを書き出す処理は次のようになってきます。

func writePacket(w io.Writer, p []byte) error {
    p = encrypt(p)
    if err := binary.Write(w, binary.LittleEndian, uint16(len(p))); err != nil {
        return err
    }
    if _, err := w.Write(p); err != nil {
        return err
    }
    return nil
}

さらに、「暗号化するかどうかはON/OFFできるようにしてほしい！」という追加の仕様がきたら、暗号化をするかどうかのフラグも追加が必要となり、だんだん writePacket 関数が複雑化していきます・・・。

func writePacket(w io.Writer, p []byte, useEncryption bool) error {
    if (useEncryption) {
        p := encrypt(p)
    }
    if err := binary.Write(w, binary.LittleEndian, uint16(len(p))); err != nil {
        return err
    }
    if _, err := w.Write(p); err != nil {
        return err
    }
    return nil
}

x/text/transformer でスマートな設計に

パケットを読み書きする関数が複雑化していくのはできれば避けたいところです。
もっとスマートな設計にできないでしょうか？

Goにはx/text/transformというパッケージが用意されています。
その中にあるTransformerインターフェースが今回使う重要なものとなります。

Transformerは簡潔に言うと、入力した[]byteに何らかの変換処理をかけるものとなります。
さらに、複数のTransformerをつないで1つのTransformerにできます。

このインターフェースを利用して次の2つのTransformerを作ると要件が実現できそうです。

先頭に長さをつけるTransformer (PacketWriteTransformer)
暗号化するTransformer (EncryptTransformer)

type PacketWriteTransformer struct{ transform.NopResetter }

func (p *PacketWriteTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    nSrc = len(src)
    binary.LittleEndian.PutUint16(dst, uint16(nSrc))
    nDst = copy(dst[2:], src)
    if nDst < nSrc {
        err = transform.ErrShortDst
    }
    nDst += 2
    return
}

type EncryptTransformer struct { transform.NopResetter }

func (p *EncryptTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {   
    encrypt(dst, src)
    nSrc = len(src)
    nDst = len(dst)
    return
}

では、この2つのTransformerを使ってPayloadが暗号化されたパケットを書き出してみます！

func main() {
    var b bytes.Buffer
    payload := []byte("hello, world")
    w := transform.NewWriter(&b, transform.Chain(&EncryptTransformer{}, &PacketWriteTransformer{})
    w.Write(&b, payload)
    print(hex.Dump(b.Bytes())) // 暗号化されたパケットになる
}

transform.Chainで2つのTransformerをつないでいるところがポイントです。
もし暗号化をOFFにしたくなったら、Chainの中からEncryptTransformerを外すだけでOKです。スマートですね！

読み出しのTransformerも作る

ここまでで、パケットの書き出し（Payload単体からパケットを作る）処理をTransformerで実現できました。
あとは、逆の処理をするTransformerを作れば、パケットの読み出しも実現できます。

先頭に長さがついたパケットからPayloadを取り出すTransformer (PacketReadTransformer)
復号化するTransformer (DecryptTransformer)

実際に書いてみると次のようになりました。

type PacketReadTransformer struct {
    rest []byte
}

func (p *PacketReadTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
    npRest := len(p.rest)
    if npRest > 0 {
        src = append(p.rest, src...)
        p.rest = nil
    }
    if atEOF && len(src) == 0 {
        return
    }
    if len(src) < 2 {
        err = transform.ErrShortSrc
        return
    }
    size := int(binary.LittleEndian.Uint16(src[:2]))
    if len(src[2:]) < size {
        err = transform.ErrShortSrc
        return
    }
    nDst = copy(dst, src[2:2+size])
    nSrc = 2 + nDst - npRest
    if nDst < size {
        err = transform.ErrShortDst
        return
    }
    nRest := len(src[2+size:])
    if nRest > 0 {
        p.rest = make([]byte, nRest)
        nSrc += copy(p.rest, src[2+size:])
    }
    return
}

func (p *PacketReadTransformer) Reset() {
    p.rest = nil
}

type DecryptTransformer struct { transform.NopResetter }

func (p *DecryptTransformer) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {   
    decrypt(dst, src)
    nSrc = len(src)
    nDst = len(dst)
    return
}

PacketReadTransformerは、書き出し側と比べるとちょっと複雑です。
なぜなら、TCPなどのストリーム指向でデータを読み出している場合、複数のパケットが連結された状態で受け取る可能性もあるからです。

後続のパケットがくっついている場合は、後続分はp.restという場所に保持しておいて、次のTransformのときに先頭にくっつけることで対応しています。

これで、パケット読み出しのTransformerもできました！
受け取ったパケットの中身を取り出す処理は次のようなコードになります。

func main() {
    var b bytes.Buffer
    packet := []byte{ ... } // 受け取ったパケット
    r := transform.NewReader(&b, transform.Chain(&PacketReadTransformer{}, &DecryptTransformer{})
    payload := make([]byte, 1024)
    n, _ := r.Read(&b, payload)
    print(hex.Dump(payload[:n]) // 受け取ったパケットの中身がみえる
}

io.Reader/io.Writerとの連携

以上で、パケットの読み書きがすべてTransformerで実現できました。
図で表すと次のようなデータの流れになります。

transform.Transformerを使う最大の利点は、transformer.Reader/transformer.Writerがio.Reader/io.Writerを実装していることです。

io.Readerはシンプルながらかなりの応用が効くGoらしいインターフェースとなっています。詳しい利点や応用例については次の記事などを読むとよいでしょう。

とにかく、Transformerを組み合わせることで、間にどれだけ処理を挟んでも、Go標準のio.Reader/io.Writerで読み書きできるため、TCPソケットやUDPソケット(net.Conn)にそのまま接続することも可能です。

すばらしいですね！