forked from anacrolix/torrent
-
Notifications
You must be signed in to change notification settings - Fork 0
/
encryption.go
62 lines (51 loc) · 1.76 KB
/
encryption.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
package btprotocol
import (
"bytes"
"io"
"github.com/james-lawrence/torrent/mse"
"github.com/pkg/errors"
)
// EncryptionHandshake encrypt a net.Conn
type EncryptionHandshake struct {
Keys mse.SecretKeyIter
mse.CryptoSelector // available crypto algorithms
}
// Incoming establishes an encrypted connection. if it is unable to establish an encrypted connection.
// it returns a buffed read/writer that contains the previously read data from the input.
// when the encryption handshake is successful, buffed === updated.
func (t EncryptionHandshake) Incoming(rw io.ReadWriter) (updated io.ReadWriter, buffed io.ReadWriter, err error) {
type buffered struct {
io.Reader
io.Writer
}
var (
buf = bytes.NewBuffer(make([]byte, 0, 68))
)
// read the bittorrent protocol magic string initially
// in case this connection is not an encrypted connection.
if _, err := io.CopyN(buf, rw, 20); err != nil {
return updated, updated, err
} else if buf.String() == Protocol {
return rw, buffered{
Reader: io.MultiReader(bytes.NewBuffer(buf.Bytes()), rw),
Writer: rw,
}, errors.New("unencrypted connection detected")
}
teed := io.MultiReader(bytes.NewBuffer(buf.Bytes()), rw)
buf = bytes.NewBuffer(make([]byte, 0, 68))
b := buffered{
Reader: io.TeeReader(teed, buf),
Writer: rw,
}
if updated, _, err = mse.ReceiveHandshake(b, t.Keys, t.CryptoSelector); err != nil {
return rw, buffered{
Reader: io.MultiReader(bytes.NewBuffer(buf.Bytes()), rw),
Writer: rw,
}, err
}
return updated, updated, nil
}
// Outgoing initiates an outgoing encrypted connection.
func (t EncryptionHandshake) Outgoing(rw io.ReadWriter, secret []byte, algo mse.CryptoMethod) (io.ReadWriter, mse.CryptoMethod, error) {
return mse.InitiateHandshake(rw, secret, nil, algo)
}