mse/mse.go

   1 // https://wiki.vuze.com/w/Message_Stream_Encryption
   2
   3 package mse
   4
   5 import (
   6         "bytes"
   7         "crypto/rand"
   8         "crypto/rc4"
   9         "crypto/sha1"
  10         "encoding/binary"
  11         "errors"
  12         "expvar"
  13         "fmt"
  14         "io"
  15         "io/ioutil"
  16         "math"
  17         "math/big"
  18         "strconv"
  19         "sync"
  20
  21         "github.com/bradfitz/iter"
  22 )
  23
  24 const (
  25         maxPadLen = 512
  26
  27         CryptoMethodPlaintext CryptoMethod = 1
  28         CryptoMethodRC4       CryptoMethod = 2
  29         AllSupportedCrypto                 = CryptoMethodPlaintext | CryptoMethodRC4
  30 )
  31
  32 type CryptoMethod uint32
  33
  34 var (
  35         // Prime P according to the spec, and G, the generator.
  36         p, g big.Int
  37         // The rand.Int max arg for use in newPadLen()
  38         newPadLenMax big.Int
  39         // For use in initer's hashes
  40         req1 = []byte("req1")
  41         req2 = []byte("req2")
  42         req3 = []byte("req3")
  43         // Verification constant "VC" which is all zeroes in the bittorrent
  44         // implementation.
  45         vc [8]byte
  46         // Zero padding
  47         zeroPad [512]byte
  48         // Tracks counts of received crypto_provides
  49         cryptoProvidesCount = expvar.NewMap("mseCryptoProvides")
  50 )
  51
  52 func init() {
  53         p.SetString("0xFFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A63A36210000000000090563", 0)
  54         g.SetInt64(2)
  55         newPadLenMax.SetInt64(maxPadLen + 1)
  56 }
  57
  58 func hash(parts ...[]byte) []byte {
  59         h := sha1.New()
  60         for _, p := range parts {
  61                 n, err := h.Write(p)
  62                 if err != nil {
  63                         panic(err)
  64                 }
  65                 if n != len(p) {
  66                         panic(n)
  67                 }
  68         }
  69         return h.Sum(nil)
  70 }
  71
  72 func newEncrypt(initer bool, s []byte, skey []byte) (c *rc4.Cipher) {
  73         c, err := rc4.NewCipher(hash([]byte(func() string {
  74                 if initer {
  75                         return "keyA"
  76                 } else {
  77                         return "keyB"
  78                 }
  79         }()), s, skey))
  80         if err != nil {
  81                 panic(err)
  82         }
  83         var burnSrc, burnDst [1024]byte
  84         c.XORKeyStream(burnDst[:], burnSrc[:])
  85         return
  86 }
  87
  88 type cipherReader struct {
  89         c  *rc4.Cipher
  90         r  io.Reader
  91         mu sync.Mutex
  92         be []byte
  93 }
  94
  95 func (cr *cipherReader) Read(b []byte) (n int, err error) {
  96         var be []byte
  97         cr.mu.Lock()
  98         if len(cr.be) >= len(b) {
  99                 be = cr.be
 100                 cr.be = nil
 101                 cr.mu.Unlock()
 102         } else {
 103                 cr.mu.Unlock()
 104                 be = make([]byte, len(b))
 105         }
 106         n, err = cr.r.Read(be[:len(b)])
 107         cr.c.XORKeyStream(b[:n], be[:n])
 108         cr.mu.Lock()
 109         if len(be) > len(cr.be) {
 110                 cr.be = be
 111         }
 112         cr.mu.Unlock()
 113         return
 114 }
 115
 116 func newCipherReader(c *rc4.Cipher, r io.Reader) io.Reader {
 117         return &cipherReader{c: c, r: r}
 118 }
 119
 120 type cipherWriter struct {
 121         c *rc4.Cipher
 122         w io.Writer
 123         b []byte
 124 }
 125
 126 func (cr *cipherWriter) Write(b []byte) (n int, err error) {
 127         be := func() []byte {
 128                 if len(cr.b) < len(b) {
 129                         return make([]byte, len(b))
 130                 } else {
 131                         ret := cr.b
 132                         cr.b = nil
 133                         return ret
 134                 }
 135         }()
 136         cr.c.XORKeyStream(be[:], b)
 137         n, err = cr.w.Write(be[:len(b)])
 138         if n != len(b) {
 139                 // The cipher will have advanced beyond the callers stream position.
 140                 // We can't use the cipher anymore.
 141                 cr.c = nil
 142         }
 143         if len(be) > len(cr.b) {
 144                 cr.b = be
 145         }
 146         return
 147 }
 148
 149 func newX() big.Int {
 150         var X big.Int
 151         X.SetBytes(func() []byte {
 152                 var b [20]byte
 153                 _, err := rand.Read(b[:])
 154                 if err != nil {
 155                         panic(err)
 156                 }
 157                 return b[:]
 158         }())
 159         return X
 160 }
 161
 162 func paddedLeft(b []byte, _len int) []byte {
 163         if len(b) == _len {
 164                 return b
 165         }
 166         ret := make([]byte, _len)
 167         if n := copy(ret[_len-len(b):], b); n != len(b) {
 168                 panic(n)
 169         }
 170         return ret
 171 }
 172
 173 // Calculate, and send Y, our public key.
 174 func (h *handshake) postY(x *big.Int) error {
 175         var y big.Int
 176         y.Exp(&g, x, &p)
 177         return h.postWrite(paddedLeft(y.Bytes(), 96))
 178 }
 179
 180 func (h *handshake) establishS() error {
 181         x := newX()
 182         h.postY(&x)
 183         var b [96]byte
 184         _, err := io.ReadFull(h.conn, b[:])
 185         if err != nil {
 186                 return fmt.Errorf("error reading Y: %s", err)
 187         }
 188         var Y, S big.Int
 189         Y.SetBytes(b[:])
 190         S.Exp(&Y, &x, &p)
 191         sBytes := S.Bytes()
 192         copy(h.s[96-len(sBytes):96], sBytes)
 193         return nil
 194 }
 195
 196 func newPadLen() int64 {
 197         i, err := rand.Int(rand.Reader, &newPadLenMax)
 198         if err != nil {
 199                 panic(err)
 200         }
 201         ret := i.Int64()
 202         if ret < 0 || ret > maxPadLen {
 203                 panic(ret)
 204         }
 205         return ret
 206 }
 207
 208 // Manages state for both initiating and receiving handshakes.
 209 type handshake struct {
 210         conn   io.ReadWriter
 211         s      [96]byte
 212         initer bool          // Whether we're initiating or receiving.
 213         skeys  SecretKeyIter // Skeys we'll accept if receiving.
 214         skey   []byte        // Skey we're initiating with.
 215         ia     []byte        // Initial payload. Only used by the initiator.
 216         // Return the bit for the crypto method the receiver wants to use.
 217         chooseMethod CryptoSelector
 218         // Sent to the receiver.
 219         cryptoProvides CryptoMethod
 220
 221         writeMu    sync.Mutex
 222         writes     [][]byte
 223         writeErr   error
 224         writeCond  sync.Cond
 225         writeClose bool
 226
 227         writerMu   sync.Mutex
 228         writerCond sync.Cond
 229         writerDone bool
 230 }
 231
 232 func (h *handshake) finishWriting() {
 233         h.writeMu.Lock()
 234         h.writeClose = true
 235         h.writeCond.Broadcast()
 236         h.writeMu.Unlock()
 237
 238         h.writerMu.Lock()
 239         for !h.writerDone {
 240                 h.writerCond.Wait()
 241         }
 242         h.writerMu.Unlock()
 243 }
 244
 245 func (h *handshake) writer() {
 246         defer func() {
 247                 h.writerMu.Lock()
 248                 h.writerDone = true
 249                 h.writerCond.Broadcast()
 250                 h.writerMu.Unlock()
 251         }()
 252         for {
 253                 h.writeMu.Lock()
 254                 for {
 255                         if len(h.writes) != 0 {
 256                                 break
 257                         }
 258                         if h.writeClose {
 259                                 h.writeMu.Unlock()
 260                                 return
 261                         }
 262                         h.writeCond.Wait()
 263                 }
 264                 b := h.writes[0]
 265                 h.writes = h.writes[1:]
 266                 h.writeMu.Unlock()
 267                 _, err := h.conn.Write(b)
 268                 if err != nil {
 269                         h.writeMu.Lock()
 270                         h.writeErr = err
 271                         h.writeMu.Unlock()
 272                         return
 273                 }
 274         }
 275 }
 276
 277 func (h *handshake) postWrite(b []byte) error {
 278         h.writeMu.Lock()
 279         defer h.writeMu.Unlock()
 280         if h.writeErr != nil {
 281                 return h.writeErr
 282         }
 283         h.writes = append(h.writes, b)
 284         h.writeCond.Signal()
 285         return nil
 286 }
 287
 288 func xor(dst, src []byte) (ret []byte) {
 289         max := len(dst)
 290         if max > len(src) {
 291                 max = len(src)
 292         }
 293         ret = make([]byte, 0, max)
 294         for i := range iter.N(max) {
 295                 ret = append(ret, dst[i]^src[i])
 296         }
 297         return
 298 }
 299
 300 func marshal(w io.Writer, data ...interface{}) (err error) {
 301         for _, data := range data {
 302                 err = binary.Write(w, binary.BigEndian, data)
 303                 if err != nil {
 304                         break
 305                 }
 306         }
 307         return
 308 }
 309
 310 func unmarshal(r io.Reader, data ...interface{}) (err error) {
 311         for _, data := range data {
 312                 err = binary.Read(r, binary.BigEndian, data)
 313                 if err != nil {
 314                         break
 315                 }
 316         }
 317         return
 318 }
 319
 320 // Looking for b at the end of a.
 321 func suffixMatchLen(a, b []byte) int {
 322         if len(b) > len(a) {
 323                 b = b[:len(a)]
 324         }
 325         // i is how much of b to try to match
 326         for i := len(b); i > 0; i-- {
 327                 // j is how many chars we've compared
 328                 j := 0
 329                 for ; j < i; j++ {
 330                         if b[i-1-j] != a[len(a)-1-j] {
 331                                 goto shorter
 332                         }
 333                 }
 334                 return j
 335         shorter:
 336         }
 337         return 0
 338 }
 339
 340 // Reads from r until b has been seen. Keeps the minimum amount of data in
 341 // memory.
 342 func readUntil(r io.Reader, b []byte) error {
 343         b1 := make([]byte, len(b))
 344         i := 0
 345         for {
 346                 _, err := io.ReadFull(r, b1[i:])
 347                 if err != nil {
 348                         return err
 349                 }
 350                 i = suffixMatchLen(b1, b)
 351                 if i == len(b) {
 352                         break
 353                 }
 354                 if copy(b1, b1[len(b1)-i:]) != i {
 355                         panic("wat")
 356                 }
 357         }
 358         return nil
 359 }
 360
 361 type readWriter struct {
 362         io.Reader
 363         io.Writer
 364 }
 365
 366 func (h *handshake) newEncrypt(initer bool) *rc4.Cipher {
 367         return newEncrypt(initer, h.s[:], h.skey)
 368 }
 369
 370 func (h *handshake) initerSteps() (ret io.ReadWriter, err error) {
 371         h.postWrite(hash(req1, h.s[:]))
 372         h.postWrite(xor(hash(req2, h.skey), hash(req3, h.s[:])))
 373         buf := &bytes.Buffer{}
 374         padLen := uint16(newPadLen())
 375         if len(h.ia) > math.MaxUint16 {
 376                 err = errors.New("initial payload too large")
 377                 return
 378         }
 379         err = marshal(buf, vc[:], h.cryptoProvides, padLen, zeroPad[:padLen], uint16(len(h.ia)), h.ia)
 380         if err != nil {
 381                 return
 382         }
 383         e := h.newEncrypt(true)
 384         be := make([]byte, buf.Len())
 385         e.XORKeyStream(be, buf.Bytes())
 386         h.postWrite(be)
 387         bC := h.newEncrypt(false)
 388         var eVC [8]byte
 389         bC.XORKeyStream(eVC[:], vc[:])
 390         // Read until the all zero VC. At this point we've only read the 96 byte
 391         // public key, Y. There is potentially 512 byte padding, between us and
 392         // the 8 byte verification constant.
 393         err = readUntil(io.LimitReader(h.conn, 520), eVC[:])
 394         if err != nil {
 395                 if err == io.EOF {
 396                         err = errors.New("failed to synchronize on VC")
 397                 } else {
 398                         err = fmt.Errorf("error reading until VC: %s", err)
 399                 }
 400                 return
 401         }
 402         r := newCipherReader(bC, h.conn)
 403         var method CryptoMethod
 404         err = unmarshal(r, &method, &padLen)
 405         if err != nil {
 406                 return
 407         }
 408         _, err = io.CopyN(ioutil.Discard, r, int64(padLen))
 409         if err != nil {
 410                 return
 411         }
 412         switch method & h.cryptoProvides {
 413         case CryptoMethodRC4:
 414                 ret = readWriter{r, &cipherWriter{e, h.conn, nil}}
 415         case CryptoMethodPlaintext:
 416                 ret = h.conn
 417         default:
 418                 err = fmt.Errorf("receiver chose unsupported method: %x", method)
 419         }
 420         return
 421 }
 422
 423 var ErrNoSecretKeyMatch = errors.New("no skey matched")
 424
 425 func (h *handshake) receiverSteps() (ret io.ReadWriter, err error) {
 426         // There is up to 512 bytes of padding, then the 20 byte hash.
 427         err = readUntil(io.LimitReader(h.conn, 532), hash(req1, h.s[:]))
 428         if err != nil {
 429                 if err == io.EOF {
 430                         err = errors.New("failed to synchronize on S hash")
 431                 }
 432                 return
 433         }
 434         var b [20]byte
 435         _, err = io.ReadFull(h.conn, b[:])
 436         if err != nil {
 437                 return
 438         }
 439         err = ErrNoSecretKeyMatch
 440         h.skeys(func(skey []byte) bool {
 441                 if bytes.Equal(xor(hash(req2, skey), hash(req3, h.s[:])), b[:]) {
 442                         h.skey = skey
 443                         err = nil
 444                         return false
 445                 }
 446                 return true
 447         })
 448         if err != nil {
 449                 return
 450         }
 451         r := newCipherReader(newEncrypt(true, h.s[:], h.skey), h.conn)
 452         var (
 453                 vc       [8]byte
 454                 provides CryptoMethod
 455                 padLen   uint16
 456         )
 457
 458         err = unmarshal(r, vc[:], &provides, &padLen)
 459         if err != nil {
 460                 return
 461         }
 462         cryptoProvidesCount.Add(strconv.FormatUint(uint64(provides), 16), 1)
 463         chosen := h.chooseMethod(provides)
 464         _, err = io.CopyN(ioutil.Discard, r, int64(padLen))
 465         if err != nil {
 466                 return
 467         }
 468         var lenIA uint16
 469         unmarshal(r, &lenIA)
 470         if lenIA != 0 {
 471                 h.ia = make([]byte, lenIA)
 472                 unmarshal(r, h.ia)
 473         }
 474         buf := &bytes.Buffer{}
 475         w := cipherWriter{h.newEncrypt(false), buf, nil}
 476         padLen = uint16(newPadLen())
 477         err = marshal(&w, &vc, uint32(chosen), padLen, zeroPad[:padLen])
 478         if err != nil {
 479                 return
 480         }
 481         err = h.postWrite(buf.Bytes())
 482         if err != nil {
 483                 return
 484         }
 485         switch chosen {
 486         case CryptoMethodRC4:
 487                 ret = readWriter{
 488                         io.MultiReader(bytes.NewReader(h.ia), r),
 489                         &cipherWriter{w.c, h.conn, nil},
 490                 }
 491         case CryptoMethodPlaintext:
 492                 ret = readWriter{
 493                         io.MultiReader(bytes.NewReader(h.ia), h.conn),
 494                         h.conn,
 495                 }
 496         default:
 497                 err = errors.New("chosen crypto method is not supported")
 498         }
 499         return
 500 }
 501
 502 func (h *handshake) Do() (ret io.ReadWriter, err error) {
 503         h.writeCond.L = &h.writeMu
 504         h.writerCond.L = &h.writerMu
 505         go h.writer()
 506         defer func() {
 507                 h.finishWriting()
 508                 if err == nil {
 509                         err = h.writeErr
 510                 }
 511         }()
 512         err = h.establishS()
 513         if err != nil {
 514                 err = fmt.Errorf("error while establishing secret: %s", err)
 515                 return
 516         }
 517         pad := make([]byte, newPadLen())
 518         io.ReadFull(rand.Reader, pad)
 519         err = h.postWrite(pad)
 520         if err != nil {
 521                 return
 522         }
 523         if h.initer {
 524                 ret, err = h.initerSteps()
 525         } else {
 526                 ret, err = h.receiverSteps()
 527         }
 528         return
 529 }
 530
 531 func InitiateHandshake(rw io.ReadWriter, skey []byte, initialPayload []byte, cryptoProvides CryptoMethod) (ret io.ReadWriter, err error) {
 532         h := handshake{
 533                 conn:           rw,
 534                 initer:         true,
 535                 skey:           skey,
 536                 ia:             initialPayload,
 537                 cryptoProvides: cryptoProvides,
 538         }
 539         return h.Do()
 540 }
 541
 542 func ReceiveHandshake(rw io.ReadWriter, skeys SecretKeyIter, selectCrypto CryptoSelector) (ret io.ReadWriter, err error) {
 543         h := handshake{
 544                 conn:         rw,
 545                 initer:       false,
 546                 skeys:        skeys,
 547                 chooseMethod: selectCrypto,
 548         }
 549         return h.Do()
 550 }
 551
 552 // A function that given a function, calls it with secret keys until it
 553 // returns false or exhausted.
 554 type SecretKeyIter func(callback func(skey []byte) (more bool))
 555
 556 func DefaultCryptoSelector(provided CryptoMethod) CryptoMethod {
 557         if provided&CryptoMethodPlaintext != 0 {
 558                 return CryptoMethodPlaintext
 559         }
 560         return CryptoMethodRC4
 561 }
 562
 563 type CryptoSelector func(CryptoMethod) CryptoMethod