83 lines
1.7 KiB
Go
83 lines
1.7 KiB
Go
|
package secure
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"crypto/rand"
|
||
|
|
"errors"
|
||
|
|
"net"
|
||
|
|
|
||
|
|
"golang.org/x/crypto/nacl/box"
|
||
|
|
)
|
||
|
|
|
||
|
|
type SecureMessage struct {
|
||
|
|
Msg []byte
|
||
|
|
Nonce [24]byte
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *SecureMessage) toByteArray() []byte {
|
||
|
|
return append(s.Nonce[:], s.Msg[:]...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func ConstructSecureMessage(sm []byte) SecureMessage {
|
||
|
|
var nonce [24]byte
|
||
|
|
nonceArray := sm[:24]
|
||
|
|
copy(nonce[:], nonceArray)
|
||
|
|
|
||
|
|
// Trim out all unnecessary bytes
|
||
|
|
msg := bytes.Trim(sm[24:], "\x00")
|
||
|
|
|
||
|
|
return SecureMessage{Msg: msg, Nonce: nonce}
|
||
|
|
}
|
||
|
|
|
||
|
|
type SecureConnection struct {
|
||
|
|
Conn *net.TCPConn
|
||
|
|
SharedKey *[32]byte
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *SecureConnection) Read(p []byte) (int, error) {
|
||
|
|
message := make([]byte, 2048)
|
||
|
|
// Read the message from the buffer
|
||
|
|
n, err := s.Conn.Read(message)
|
||
|
|
|
||
|
|
secureMessage := ConstructSecureMessage(message)
|
||
|
|
decryptedMessage, ok := box.OpenAfterPrecomputation(nil, secureMessage.Msg, &secureMessage.Nonce, s.SharedKey)
|
||
|
|
|
||
|
|
if !ok {
|
||
|
|
return 0, errors.New("Problem decrypting the message.\n")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Actually copy it to the destination byte array
|
||
|
|
n = copy(p, decryptedMessage)
|
||
|
|
|
||
|
|
return n, err
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *SecureConnection) Write(p []byte) (int, error) {
|
||
|
|
var nonce [24]byte
|
||
|
|
|
||
|
|
// Create a new nonce for each message sent
|
||
|
|
rand.Read(nonce[:])
|
||
|
|
|
||
|
|
encryptedMessage := box.SealAfterPrecomputation(nil, p, &nonce, s.SharedKey)
|
||
|
|
sm := SecureMessage{Msg: encryptedMessage, Nonce: nonce}
|
||
|
|
|
||
|
|
// Write it to the connection
|
||
|
|
return s.Conn.Write(sm.toByteArray())
|
||
|
|
}
|
||
|
|
|
||
|
|
func Handshake(conn *net.TCPConn) *[32]byte {
|
||
|
|
var peerKey, sharedKey [32]byte
|
||
|
|
|
||
|
|
publicKey, privateKey, _ := box.GenerateKey(rand.Reader)
|
||
|
|
|
||
|
|
conn.Write(publicKey[:])
|
||
|
|
|
||
|
|
peerKeyArray := make([]byte, 32)
|
||
|
|
conn.Read(peerKeyArray)
|
||
|
|
copy(peerKey[:], peerKeyArray)
|
||
|
|
|
||
|
|
box.Precompute(&sharedKey, &peerKey, privateKey)
|
||
|
|
|
||
|
|
return &sharedKey
|
||
|
|
}
|