secio -- spipe v2

This commit introduces secio, the next generation spipe.

crypto/secio/al.go

@@ -0,0 +1,116 @@
+package secio
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/hmac"
+	"crypto/sha1"
+	"crypto/sha256"
+	"crypto/sha512"
+	"hash"
+
+	bfish "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.crypto/blowfish"
+
+	ci "github.com/jbenet/go-ipfs/crypto"
+)
+
+// List of supported ECDH curves
+var SupportedExchanges = "P-256,P-224,P-384,P-521"
+
+// List of supported Ciphers
+var SupportedCiphers = "AES-256,AES-128,Blowfish"
+
+// List of supported Hashes
+var SupportedHashes = "SHA256,SHA512"
+
+type HMAC struct {
+	hash.Hash
+	size int
+}
+
+// encParams represent encryption parameters
+type encParams struct {
+	// keys
+	permanentPubKey ci.PubKey
+	ephemeralPubKey []byte
+	keys            ci.StretchedKeys
+
+	// selections
+	curveT  string
+	cipherT string
+	hashT   string
+
+	// cipher + mac
+	cipher cipher.Stream
+	mac    HMAC
+}
+
+func (e *encParams) makeMacAndCipher() error {
+	m, err := newMac(e.hashT, e.keys.MacKey)
+	if err != nil {
+		return err
+	}
+
+	bc, err := newBlockCipher(e.cipherT, e.keys.CipherKey)
+	if err != nil {
+		return err
+	}
+
+	e.cipher = cipher.NewCTR(bc, e.keys.IV)
+	e.mac = m
+	return nil
+}
+
+func newMac(hashType string, key []byte) (HMAC, error) {
+	switch hashType {
+	case "SHA1":
+		return HMAC{hmac.New(sha1.New, key), sha1.Size}, nil
+	case "SHA512":
+		return HMAC{hmac.New(sha512.New, key), sha512.Size}, nil
+	case "SHA256":
+		return HMAC{hmac.New(sha256.New, key), sha256.Size}, nil
+	default:
+		return HMAC{}, fmt.Errorf("Unrecognized hash type: %s", hashType)
+	}
+}
+
+func newBlockCipher(cipherT string, key []byte) (cipher.Block, error) {
+	switch cipherT {
+	case "AES-128", "AES-256":
+		return aes.NewCipher(key)
+	case "Blowfish":
+		return bfish.NewCipher(key)
+	default:
+		return nil, fmt.Errorf("Unrecognized cipher type: %s", cipherT)
+	}
+}
+
+// Determines which algorithm to use.  Note:  f(a, b) = f(b, a)
+func selectBest(order int, p1, p2 string) (string, error) {
+	var f, s []string
+	switch order {
+	case -1:
+		f = strings.Split(p2, ",")
+		s = strings.Split(p1, ",")
+	case 1:
+		f = strings.Split(p1, ",")
+		s = strings.Split(p2, ",")
+	default: // Exact same preferences.
+		p := strings.Split(p1, ",")
+		return p[0], nil
+	}
+
+	for _, fc := range f {
+		for _, sc := range s {
+			if fc == sc {
+				return fc, nil
+			}
+		}
+	}
+
+	return "", errors.New("No algorithms in common!")
+}

crypto/secio/interface.go

@@ -0,0 +1,72 @@
+// package secio handles establishing secure communication between two peers.
+package secio
+
+import (
+	"io"
+
+	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
+	msgio "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio"
+
+	peer "github.com/jbenet/go-ipfs/peer"
+)
+
+// SessionGenerator constructs secure communication sessions for a peer.
+type SessionGenerator struct {
+	Local     peer.Peer
+	Peerstore peer.Peerstore
+}
+
+// NewSession takes an insecure io.ReadWriter, performs a TLS-like
+// handshake with the other side, and returns a secure session.
+// See the source for the protocol details and security implementation.
+// The provided Context is only needed for the duration of this function.
+func (sg *SessionGenerator) NewSession(ctx context.Context,
+	insecure io.ReadWriter) (Session, error) {
+
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	ctx, cancel := context.WithCancel(ctx)
+
+	ss := newSecureSession(sg.Local, sg.Peerstore)
+	if err := ss.handshake(ctx, insecure); err != nil {
+		cancel()
+		return nil, err
+	}
+
+	return ss, nil
+}
+
+type Session interface {
+	// ReadWriter returns the encrypted communication channel
+	ReadWriter() msgio.ReadWriteCloser
+
+	// LocalPeer retrieves the local peer.
+	LocalPeer() peer.Peer
+
+	// RemotePeer retrieves the local peer.
+	RemotePeer() peer.Peer
+
+	// Close closes the secure session
+	Close() error
+}
+
+// SecureReadWriter returns the encrypted communication channel
+func (s *secureSession) ReadWriter() msgio.ReadWriteCloser {
+	return s.secure
+}
+
+// LocalPeer retrieves the local peer.
+func (s *secureSession) LocalPeer() peer.Peer {
+	return s.localPeer
+}
+
+// RemotePeer retrieves the local peer.
+func (s *secureSession) RemotePeer() peer.Peer {
+	return s.remotePeer
+}
+
+// Close closes the secure session
+func (s *secureSession) Close() error {
+	return s.secure.Close()
+}

crypto/secio/internal/pb/Makefile

@@ -0,0 +1,10 @@
+PB = $(wildcard *.proto)
+GO = $(PB:.proto=.pb.go)
+
+all: $(GO)
+
+%.pb.go: %.proto
+		protoc --gogo_out=. --proto_path=../../../../../../:/usr/local/opt/protobuf/include:. $<
+
+clean:
+		rm *.pb.go

crypto/secio/internal/pb/spipe.pb.go

@@ -0,0 +1,99 @@
+// Code generated by protoc-gen-gogo.
+// source: spipe.proto
+// DO NOT EDIT!
+
+/*
+Package spipe_pb is a generated protocol buffer package.
+
+It is generated from these files:
+	spipe.proto
+
+It has these top-level messages:
+	Propose
+	Exchange
+*/
+package spipe_pb
+
+import proto "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/gogoprotobuf/proto"
+import json "encoding/json"
+import math "math"
+
+// Reference proto, json, and math imports to suppress error if they are not otherwise used.
+var _ = proto.Marshal
+var _ = &json.SyntaxError{}
+var _ = math.Inf
+
+type Propose struct {
+	Rand             []byte  `protobuf:"bytes,1,opt,name=rand" json:"rand,omitempty"`
+	Pubkey           []byte  `protobuf:"bytes,2,opt,name=pubkey" json:"pubkey,omitempty"`
+	Exchanges        *string `protobuf:"bytes,3,opt,name=exchanges" json:"exchanges,omitempty"`
+	Ciphers          *string `protobuf:"bytes,4,opt,name=ciphers" json:"ciphers,omitempty"`
+	Hashes           *string `protobuf:"bytes,5,opt,name=hashes" json:"hashes,omitempty"`
+	XXX_unrecognized []byte  `json:"-"`
+}
+
+func (m *Propose) Reset()         { *m = Propose{} }
+func (m *Propose) String() string { return proto.CompactTextString(m) }
+func (*Propose) ProtoMessage()    {}
+
+func (m *Propose) GetRand() []byte {
+	if m != nil {
+		return m.Rand
+	}
+	return nil
+}
+
+func (m *Propose) GetPubkey() []byte {
+	if m != nil {
+		return m.Pubkey
+	}
+	return nil
+}
+
+func (m *Propose) GetExchanges() string {
+	if m != nil && m.Exchanges != nil {
+		return *m.Exchanges
+	}
+	return ""
+}
+
+func (m *Propose) GetCiphers() string {
+	if m != nil && m.Ciphers != nil {
+		return *m.Ciphers
+	}
+	return ""
+}
+
+func (m *Propose) GetHashes() string {
+	if m != nil && m.Hashes != nil {
+		return *m.Hashes
+	}
+	return ""
+}
+
+type Exchange struct {
+	Epubkey          []byte `protobuf:"bytes,1,opt,name=epubkey" json:"epubkey,omitempty"`
+	Signature        []byte `protobuf:"bytes,2,opt,name=signature" json:"signature,omitempty"`
+	XXX_unrecognized []byte `json:"-"`
+}
+
+func (m *Exchange) Reset()         { *m = Exchange{} }
+func (m *Exchange) String() string { return proto.CompactTextString(m) }
+func (*Exchange) ProtoMessage()    {}
+
+func (m *Exchange) GetEpubkey() []byte {
+	if m != nil {
+		return m.Epubkey
+	}
+	return nil
+}
+
+func (m *Exchange) GetSignature() []byte {
+	if m != nil {
+		return m.Signature
+	}
+	return nil
+}
+
+func init() {
+}

crypto/secio/internal/pb/spipe.proto

@@ -0,0 +1,14 @@
+package spipe.pb;
+
+message Propose {
+	optional bytes rand = 1;
+	optional bytes pubkey = 2;
+	optional string exchanges = 3;
+	optional string ciphers = 4;
+	optional string hashes = 5;
+}
+
+message Exchange {
+	optional bytes epubkey = 1;
+	optional bytes signature = 2;
+}

crypto/secio/io_test.go

@@ -0,0 +1 @@
+package secio

crypto/secio/protocol.go

@@ -0,0 +1,282 @@
+package secio
+
+import (
+	"bytes"
+	"crypto/rand"
+	"errors"
+	"fmt"
+	"io"
+
+	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
+	msgio "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio"
+
+	ci "github.com/jbenet/go-ipfs/crypto"
+	pb "github.com/jbenet/go-ipfs/crypto/spipe/internal/pb"
+	peer "github.com/jbenet/go-ipfs/peer"
+	u "github.com/jbenet/go-ipfs/util"
+)
+
+var log = u.Logger("spipe")
+
+// ErrUnsupportedKeyType is returned when a private key cast/type switch fails.
+var ErrUnsupportedKeyType = errors.New("unsupported key type")
+
+// ErrClosed signals the closing of a connection.
+var ErrClosed = errors.New("connection closed")
+
+// nonceSize is the size of our nonces (in bytes)
+const nonceSize = 16
+
+// secureSession encapsulates all the parameters needed for encrypting
+// and decrypting traffic from an insecure channel.
+type secureSession struct {
+	secure msgio.ReadWriteCloser
+
+	insecure  io.ReadWriter
+	insecureM msgio.ReadWriter
+
+	peers      peer.Peerstore
+	localPeer  peer.Peer
+	remotePeer peer.Peer
+
+	local  encParams
+	remote encParams
+
+	sharedSecret []byte
+}
+
+func newSecureSession(local peer.Peer, peers peer.Peerstore) *secureSession {
+	return &secureSession{peers: peers, localPeer: local}
+}
+
+// handsahke performs initial communication over insecure channel to share
+// keys, IDs, and initiate communication, assigning all necessary params.
+// requires the duplex channel to be a msgio.ReadWriter (for framed messaging)
+func (s *secureSession) handshake(ctx context.Context, insecure io.ReadWriter) error {
+
+	s.insecure = insecure
+	s.insecureM = msgio.NewReadWriter(insecure)
+
+	// =============================================================================
+	// step 1. Propose -- propose cipher suite + send pubkeys + nonce
+
+	// Generate and send Hello packet.
+	// Hello = (rand, PublicKey, Supported)
+	nonceOut := make([]byte, nonceSize)
+	_, err := rand.Read(nonceOut)
+	if err != nil {
+		return err
+	}
+
+	log.Debugf("handshake: %s <--start--> %s", s.localPeer, s.remotePeer)
+	s.local.permanentPubKey = s.localPeer.PubKey()
+	myPubKeyBytes, err := s.local.permanentPubKey.Bytes()
+	if err != nil {
+		return err
+	}
+
+	proposeOut := new(pb.Propose)
+	proposeOut.Rand = nonceOut
+	proposeOut.Pubkey = myPubKeyBytes
+	proposeOut.Exchanges = &SupportedExchanges
+	proposeOut.Ciphers = &SupportedCiphers
+	proposeOut.Hashes = &SupportedHashes
+
+	// Send Propose packet (respects ctx)
+	proposeOutBytes, err := writeMsgCtx(ctx, s.insecureM, proposeOut)
+	if err != nil {
+		return err
+	}
+
+	// Receive + Parse their Propose packet and generate an Exchange packet.
+	proposeIn := new(pb.Propose)
+	proposeInBytes, err := readMsgCtx(ctx, s.insecureM, proposeIn)
+	if err != nil {
+		return err
+	}
+
+	// =============================================================================
+	// step 1.1 Identify -- get identity from their key
+
+	// get remote identity
+	s.remote.permanentPubKey, err = ci.UnmarshalPublicKey(proposeIn.GetPubkey())
+	if err != nil {
+		return err
+	}
+
+	// get or construct peer
+	s.remotePeer, err = getOrConstructPeer(s.peers, s.remote.permanentPubKey)
+	if err != nil {
+		return err
+	}
+	// log.Debugf("%s Remote Peer Identified as %s", s.localPeer, s.remotePeer)
+
+	// =============================================================================
+	// step 1.2 Selection -- select/agree on best encryption parameters
+
+	// to determine order, use cmp(H(lr||rpk), H(rr||lpk)).
+	oh1 := u.Hash(append(proposeIn.GetPubkey(), nonceOut...))
+	oh2 := u.Hash(append(myPubKeyBytes, proposeIn.GetRand()...))
+	order := bytes.Compare(oh1, oh2)
+	s.local.curveT, err = selectBest(order, SupportedExchanges, proposeIn.GetExchanges())
+	if err != nil {
+		return err
+	}
+
+	s.local.cipherT, err = selectBest(order, SupportedCiphers, proposeIn.GetCiphers())
+	if err != nil {
+		return err
+	}
+
+	s.local.hashT, err = selectBest(order, SupportedHashes, proposeIn.GetHashes())
+	if err != nil {
+		return err
+	}
+
+	// we use the same params for both directions (must choose same curve)
+	// WARNING: if they dont SelectBest the same way, this won't work...
+	s.remote.curveT = s.local.curveT
+	s.remote.cipherT = s.local.cipherT
+	s.remote.hashT = s.local.hashT
+
+	// =============================================================================
+	// step 2. Exchange -- exchange (signed) ephemeral keys. verify signatures.
+
+	// Generate EphemeralPubKey
+	var genSharedKey ci.GenSharedKey
+	s.local.ephemeralPubKey, genSharedKey, err = ci.GenerateEKeyPair(s.local.curveT)
+
+	// Gather corpus to sign.
+	var selectionOut bytes.Buffer
+	selectionOut.Write(proposeOutBytes)
+	selectionOut.Write(proposeInBytes)
+	selectionOut.Write(s.local.ephemeralPubKey)
+	selectionOutBytes := selectionOut.Bytes()
+
+	exchangeOut := new(pb.Exchange)
+	exchangeOut.Epubkey = s.local.ephemeralPubKey
+	exchangeOut.Signature, err = s.localPeer.PrivKey().Sign(selectionOutBytes)
+	if err != nil {
+		return err
+	}
+
+	// Send Propose packet (respects ctx)
+	if _, err := writeMsgCtx(ctx, s.insecureM, exchangeOut); err != nil {
+		return err
+	}
+
+	// Receive + Parse their Propose packet and generate an Exchange packet.
+	exchangeIn := new(pb.Exchange)
+	if _, err := readMsgCtx(ctx, s.insecureM, exchangeIn); err != nil {
+		return err
+	}
+
+	// =============================================================================
+	// step 2.1. Verify -- verify their exchange packet is good.
+
+	// get their ephemeral pub key
+	s.remote.ephemeralPubKey = exchangeIn.GetEpubkey()
+
+	var selectionIn bytes.Buffer
+	selectionIn.Write(proposeInBytes)
+	selectionIn.Write(proposeOutBytes)
+	selectionIn.Write(s.remote.ephemeralPubKey)
+	selectionInBytes := selectionIn.Bytes()
+
+	// u.POut("Remote Peer Identified as %s\n", s.remote)
+	sigOK, err := s.remotePeer.PubKey().Verify(selectionInBytes, exchangeIn.GetSignature())
+	if err != nil {
+		return err
+	}
+
+	if !sigOK {
+		return errors.New("Bad signature!")
+	}
+
+	// =============================================================================
+	// step 2.2. Keys -- generate keys for mac + encryption
+
+	// OK! seems like we're good to go.
+	s.sharedSecret, err = genSharedKey(exchangeIn.GetEpubkey())
+	if err != nil {
+		return err
+	}
+
+	// generate two sets of keys (stretching)
+	k1, k2 := ci.KeyStretcher(s.local.cipherT, s.local.hashT, s.sharedSecret)
+
+	// use random nonces to decide order.
+	switch order {
+	case 1:
+	case -1:
+		k1, k2 = k2, k1 // swap
+	default:
+		log.Error("WOAH: same keys (AND same nonce: 1/(2^128) chance!).")
+		// this shouldn't happen. must determine order another way.
+		// use the same keys but, make sure to copy underlying data!
+		copy(k2.IV, k1.IV)
+		copy(k2.MacKey, k1.MacKey)
+		copy(k2.CipherKey, k1.CipherKey)
+	}
+	s.local.keys = k1
+	s.remote.keys = k2
+
+	// =============================================================================
+	// step 2.3. MAC + Cipher -- prepare MAC + cipher
+
+	if err := s.local.makeMacAndCipher(); err != nil {
+		return err
+	}
+
+	if err := s.remote.makeMacAndCipher(); err != nil {
+		return err
+	}
+
+	// =============================================================================
+	// step 3. Finish -- send expected message (the nonces), verify encryption works
+
+	// setup ETM ReadWriter
+	w := NewETMWriter(s.insecure, s.local.cipher, s.local.mac)
+	r := NewETMReader(s.insecure, s.remote.cipher, s.remote.mac)
+	s.secure = msgio.Combine(w, r).(msgio.ReadWriteCloser)
+
+	// send their Nonce.
+	if _, err := s.secure.Write(proposeIn.GetRand()); err != nil {
+		return fmt.Errorf("Failed to write Finish nonce: %s", err)
+	}
+
+	// read our Nonce
+	nonceOut2 := make([]byte, len(nonceOut))
+	if _, err := io.ReadFull(s.secure, nonceOut2); err != nil {
+		return fmt.Errorf("Failed to read Finish nonce: %s", err)
+	}
+	if !bytes.Equal(nonceOut, nonceOut2) {
+		return fmt.Errorf("Failed to read our encrypted nonce, go: %s", nonceOut2)
+	}
+
+	// Whew! ok, that's all folks.
+	log.Debugf("handshake: %s <--finish--> %s", s.localPeer, s.remotePeer)
+	return nil
+}
+
+// getOrConstructPeer attempts to fetch a peer from a peerstore.
+// if succeeds, verify ID and PubKey match.
+// else, construct it.
+func getOrConstructPeer(peers peer.Peerstore, rpk ci.PubKey) (peer.Peer, error) {
+
+	rid, err := peer.IDFromPubKey(rpk)
+	if err != nil {
+		return nil, err
+	}
+
+	npeer, err := peers.FindOrCreate(rid)
+	if err != nil {
+		return nil, err // unexpected error happened.
+	}
+
+	// public key verification happens in Peer.VerifyAndSetPubKey
+	if err := npeer.VerifyAndSetPubKey(rpk); err != nil {
+		return nil, err // pubkey mismatch or other problem
+	}
+	return npeer, nil
+}

crypto/secio/rw.go

@@ -0,0 +1,197 @@
+package secio
+
+import (
+	"crypto/cipher"
+	"errors"
+	"fmt"
+	"io"
+
+	"crypto/hmac"
+
+	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
+	proto "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/goprotobuf/proto"
+	msgio "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio"
+)
+
+// ErrMACInvalid signals that a MAC verification failed
+var ErrMACInvalid = errors.New("MAC verification failed")
+
+type etmWriter struct {
+	// params
+	msg msgio.WriteCloser
+	str cipher.Stream
+	mac HMAC
+}
+
+// NewETMWriter Encrypt-Then-MAC
+func NewETMWriter(w io.Writer, s cipher.Stream, mac HMAC) msgio.WriteCloser {
+	return &etmWriter{msg: msgio.NewWriter(w), str: s, mac: mac}
+}
+
+// Write writes passed in buffer as a single message.
+func (w *etmWriter) Write(b []byte) (int, error) {
+	if err := w.WriteMsg(b); err != nil {
+		return 0, err
+	}
+	return len(b), nil
+}
+
+// WriteMsg writes the msg in the passed in buffer.
+func (w *etmWriter) WriteMsg(b []byte) error {
+
+	// encrypt.
+	w.str.XORKeyStream(b, b)
+
+	// then, mac.
+	if _, err := w.mac.Write(b); err != nil {
+		return err
+	}
+
+	// Sum appends.
+	b = w.mac.Sum(b)
+	w.mac.Reset()
+	// it's sad to append here. our buffers are -- hopefully -- coming from
+	// a shared buffer pool, so the append may not actually cause allocation
+	// one can only hope. i guess we'll see.
+
+	return w.msg.WriteMsg(b)
+}
+
+func (w *etmWriter) Close() error {
+	return w.msg.Close()
+}
+
+type etmReader struct {
+	msgio.Reader
+	io.Closer
+
+	// params
+	msg msgio.ReadCloser
+	str cipher.Stream
+	mac HMAC
+}
+
+// NewETMReader Encrypt-Then-MAC
+func NewETMReader(r io.Reader, s cipher.Stream, mac HMAC) msgio.ReadCloser {
+	return &etmReader{msg: msgio.NewReader(r), str: s, mac: mac}
+}
+
+func (r *etmReader) Read(buf []byte) (int, error) {
+	buf2 := buf
+	changed := false
+	if cap(buf2) < (len(buf) + r.mac.size) {
+		buf2 = make([]byte, len(buf)+r.mac.size)
+		changed = true
+	}
+
+	// WARNING: assumes msg.Read will only read _one_ message. this is what
+	// msgio is supposed to do. but msgio may change in the future. may this
+	// comment be your guiding light.
+	n, err := r.msg.Read(buf2)
+	if err != nil {
+		return n, err
+	}
+	buf2 = buf2[:n]
+
+	m, err := r.macCheckThenDecrypt(buf2)
+	if err != nil {
+		return 0, err
+	}
+	buf2 = buf2[:m]
+	if changed {
+		return copy(buf, buf2), nil
+	}
+	return m, nil
+}
+
+func (r *etmReader) ReadMsg() ([]byte, error) {
+	msg, err := r.msg.ReadMsg()
+	if err != nil {
+		return nil, err
+	}
+
+	n, err := r.macCheckThenDecrypt(msg)
+	if err != nil {
+		return nil, err
+	}
+	return msg[:n], nil
+}
+
+func (r *etmReader) macCheckThenDecrypt(m []byte) (int, error) {
+	l := len(m)
+	if l < r.mac.size {
+		return 0, fmt.Errorf("buffer (%d) shorter than MAC size (%d)", l, r.mac.size)
+	}
+
+	mark := l - r.mac.size
+	data := m[:mark]
+	macd := m[mark:]
+
+	r.mac.Write(data)
+	expected := r.mac.Sum(nil)
+	r.mac.Reset()
+
+	// check mac. if failed, return error.
+	if !hmac.Equal(macd, expected) {
+		log.Error("MAC Invalid:", expected, "!=", macd)
+		return 0, ErrMACInvalid
+	}
+
+	// ok seems good. decrypt.
+	r.str.XORKeyStream(data, data)
+	return mark, nil
+}
+
+func (w *etmReader) Close() error {
+	return w.msg.Close()
+}
+
+// ReleaseMsg signals a buffer can be reused.
+func (r *etmReader) ReleaseMsg(b []byte) {
+	r.msg.ReleaseMsg(b)
+}
+
+// writeMsgCtx is used by the
+func writeMsgCtx(ctx context.Context, w msgio.Writer, msg proto.Message) ([]byte, error) {
+	enc, err := proto.Marshal(msg)
+	if err != nil {
+		return nil, err
+	}
+
+	// write in a goroutine so we can exit when our context is cancelled.
+	done := make(chan error)
+	go func(m []byte) {
+		err := w.WriteMsg(m)
+		done <- err
+	}(enc)
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case e := <-done:
+		return enc, e
+	}
+}
+
+func readMsgCtx(ctx context.Context, r msgio.Reader, p proto.Message) ([]byte, error) {
+	var msg []byte
+
+	// read in a goroutine so we can exit when our context is cancelled.
+	done := make(chan error)
+	go func() {
+		var err error
+		msg, err = r.ReadMsg()
+		done <- err
+	}()
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case e := <-done:
+		if e != nil {
+			return nil, e
+		}
+	}
+
+	return msg, proto.Unmarshal(msg, p)
+}