forked from libp2p/go-libp2p
-
Notifications
You must be signed in to change notification settings - Fork 0
/
psk_conn.go
83 lines (69 loc) · 1.59 KB
/
psk_conn.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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
package pnet
import (
"crypto/cipher"
"crypto/rand"
"io"
"net"
"github.com/seqsy/go-libp2p/core/pnet"
"github.com/davidlazar/go-crypto/salsa20"
pool "github.com/libp2p/go-buffer-pool"
)
// we are using buffer pool as user needs their slice back
// so we can't do XOR cripter in place
var (
errShortNonce = pnet.NewError("could not read full nonce")
errInsecureNil = pnet.NewError("insecure is nil")
errPSKNil = pnet.NewError("pre-shread key is nil")
)
type pskConn struct {
net.Conn
psk *[32]byte
writeS20 cipher.Stream
readS20 cipher.Stream
}
func (c *pskConn) Read(out []byte) (int, error) {
if c.readS20 == nil {
nonce := make([]byte, 24)
_, err := io.ReadFull(c.Conn, nonce)
if err != nil {
return 0, errShortNonce
}
c.readS20 = salsa20.New(c.psk, nonce)
}
n, err := c.Conn.Read(out) // read to in
if n > 0 {
c.readS20.XORKeyStream(out[:n], out[:n]) // decrypt to out buffer
}
return n, err
}
func (c *pskConn) Write(in []byte) (int, error) {
if c.writeS20 == nil {
nonce := make([]byte, 24)
_, err := rand.Read(nonce)
if err != nil {
return 0, err
}
_, err = c.Conn.Write(nonce)
if err != nil {
return 0, err
}
c.writeS20 = salsa20.New(c.psk, nonce)
}
out := pool.Get(len(in))
defer pool.Put(out)
c.writeS20.XORKeyStream(out, in) // encrypt
return c.Conn.Write(out) // send
}
var _ net.Conn = (*pskConn)(nil)
func newPSKConn(psk *[32]byte, insecure net.Conn) (net.Conn, error) {
if insecure == nil {
return nil, errInsecureNil
}
if psk == nil {
return nil, errPSKNil
}
return &pskConn{
Conn: insecure,
psk: psk,
}, nil
}