gRPCでは基本的にSSL/TLSを用いて通信の暗号化を行うのだが、何かしらの共通鍵暗号など独自の方式でやりたい！という場合があるかもしれない。そういう場合でもうまく拡張すれば実現できる。今回はAES+CTRを用いた実装を試みたので、その紹介をする。サンプルコードはtakaishi/hello2018/tree/master/grpc_with_record_cryptにある。

どうやるかというと、net.ConnのWrite()でデータを送る際に暗号化し、Read()で受け取る際に復号する。AES+CTRを用いた暗号化はGoだとライブラリがあるのでそれを使う。Go言語と暗号技術（AESからTLS）が参考になる。

Write/Read時の暗号化/復号をどうやるかだが、TransportCredentialsのClientHandshakeとServerHandshakeがポイント。この2つの関数はnet.Conn型であるrawConnを受け取り、net.Conn型の変数を返している。つまり、ここでrawConnをラップしたnet.Conn型を満たす構造体を用意し、その構造体のWrite/Readであれこれすればよい。これはATLSが使っている手法で、実装するにあたり大変参考になった。

type TC struct {
    info   *credentials.ProtocolInfo
    secure bool
}

func (tc *TC) ClientHandshake(ctx context.Context, addr string, rawConn net.Conn) (_ net.Conn, _ credentials.AuthInfo, err error) {
    var c net.Conn
        c, err = conn.NewConn(rawConn)
    if err != nil {
        return nil, nil, err
    }
    return c, nil, err
}

func (tc *TC) ServerHandshake(rawConn net.Conn) (_ net.Conn, _ credentials.AuthInfo, err error) {
    var c net.Conn
    c, err = conn.NewConn(rawConn)
    if err != nil {
        return nil, nil, err
    }
    return c, nil, err
}

net.Connをラップしたconnの例。これは暗号化・複合の処理は行わず、単純にラップするだけ。

type conn struct {
    net.Conn
}

func (p *conn) Read(b []byte) (int, error) {
    return p.Conn.Read(b)
}

func (p *conn) Write(b []byte) (int, error) {
    return p.Conn.Write(b)
}

func NewConn(c net.Conn) (net.Conn, error) {
    conn := &conn{
        Conn: c,
    }
    return conn, nil
}

暗号化・復号する処理を実装するとこんな感じになる。暗号化されたデータは元データとサイズが変わるため、Prefixとして暗号化されたデータのサイズを先に送り、そのサイズ分だけ読んだ後に復号する、というやりかたにした。

const (
    MsgLenFieldSize = 4
)

type secureConn struct {
    net.Conn
    crypt    HelloRecordCrypt
    overhead int
}

func (p *secureConn) Read(buf []byte) (int, error) {
    var msgSize uint32
    var decrypted []byte
    var msgSizeBuf []byte
    var msgBuf []byte

    msgSizeBuf = make([]byte, p.overhead)
    _, err := p.Conn.Read(msgSizeBuf)
    if err != nil {
        return 0, err
    }
    msgSize = binary.LittleEndian.Uint32(msgSizeBuf)

    msgBuf = make([]byte, msgSize)
    _, err = p.Conn.Read(msgBuf)
    if err != nil {
        return 0, err
    }

    if msgSize != 0 {
        decrypted, err = p.crypt.Decrypt(decrypted, msgBuf)
        if err != nil {
            return 0, err
        }
    }
    n := copy(buf, decrypted)
    return n, nil
}

func (p *secureConn) Write(rawBuf []byte) (int, error) {
    var buf []byte

    buf, err := p.crypt.Encrypt(buf, rawBuf)
    if err != nil {
        return 0, err
    }
    msg := make([]byte, len(buf)+p.overhead)

    copy(msg[4:], buf)
    msgSize := uint32(len(msg) - p.overhead)
    binary.LittleEndian.PutUint32(msg, msgSize)
    _, err = p.Conn.Write(msg)
    if err != nil {
        return 0, err
    }
    return len(rawBuf), nil
}

func NewSecureConn(c net.Conn) (net.Conn, error) {
    crypt := &HelloRecordCrypt{}
    overhead := MsgLenFieldSize
    helloConn := &secureConn{
        Conn:     c,
        crypt:    *crypt,
        overhead: overhead,
    }

    return helloConn, nil
}

動かしてみる。暗号化しない場合、パケットキャプチャすると送受信したデータが読める(Hello!!!やThanks!!!など）ことがわかる。秘匿情報を送受信した場合、読まれる可能性があるだろう。

暗号化した場合はキャプチャしても以下のようになり、読むことが出来ない。共通鍵暗号なので、鍵がなければ解読できず、かなり安全になる。

ちょこちょこ書いて動くようになるまで1ヶ月ほどかかった。Read/Writeのラップがうまくいっていなかったのだけど、何もしない状態でテストを書くことで結果として暗号化・復号する場合でも同じ挙動になることが担保できたのでテストは大事だなと思う。