136 lines
3.6 KiB
Go
136 lines
3.6 KiB
Go
// Copyright (c) 2021 Tulir Asokan
|
|
//
|
|
// This Source Code Form is subject to the terms of the Mozilla Public
|
|
// License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
|
|
package socket
|
|
|
|
import (
|
|
"crypto/aes"
|
|
"crypto/cipher"
|
|
"crypto/sha256"
|
|
"fmt"
|
|
"io"
|
|
"sync/atomic"
|
|
|
|
"golang.org/x/crypto/curve25519"
|
|
"golang.org/x/crypto/hkdf"
|
|
)
|
|
|
|
type NoiseHandshake struct {
|
|
hash []byte
|
|
salt []byte
|
|
key cipher.AEAD
|
|
counter uint32
|
|
}
|
|
|
|
func NewNoiseHandshake() *NoiseHandshake {
|
|
return &NoiseHandshake{}
|
|
}
|
|
|
|
func newCipher(key []byte) (cipher.AEAD, error) {
|
|
aesCipher, err := aes.NewCipher(key)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
aesGCM, err := cipher.NewGCM(aesCipher)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return aesGCM, nil
|
|
}
|
|
|
|
func sha256Slice(data []byte) []byte {
|
|
hash := sha256.Sum256(data)
|
|
return hash[:]
|
|
}
|
|
|
|
func (nh *NoiseHandshake) Start(pattern string, header []byte) {
|
|
data := []byte(pattern)
|
|
if len(data) == 32 {
|
|
nh.hash = data
|
|
} else {
|
|
nh.hash = sha256Slice(data)
|
|
}
|
|
nh.salt = nh.hash
|
|
var err error
|
|
nh.key, err = newCipher(nh.hash)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
nh.Authenticate(header)
|
|
}
|
|
|
|
func (nh *NoiseHandshake) Authenticate(data []byte) {
|
|
nh.hash = sha256Slice(append(nh.hash, data...))
|
|
}
|
|
|
|
func (nh *NoiseHandshake) postIncrementCounter() uint32 {
|
|
count := atomic.AddUint32(&nh.counter, 1)
|
|
return count - 1
|
|
}
|
|
|
|
func (nh *NoiseHandshake) Encrypt(plaintext []byte) []byte {
|
|
ciphertext := nh.key.Seal(nil, generateIV(nh.postIncrementCounter()), plaintext, nh.hash)
|
|
nh.Authenticate(ciphertext)
|
|
return ciphertext
|
|
}
|
|
|
|
func (nh *NoiseHandshake) Decrypt(ciphertext []byte) (plaintext []byte, err error) {
|
|
plaintext, err = nh.key.Open(nil, generateIV(nh.postIncrementCounter()), ciphertext, nh.hash)
|
|
if err == nil {
|
|
nh.Authenticate(ciphertext)
|
|
}
|
|
return
|
|
}
|
|
|
|
func (nh *NoiseHandshake) Finish(fs *FrameSocket, frameHandler FrameHandler, disconnectHandler DisconnectHandler) (*NoiseSocket, error) {
|
|
if write, read, err := nh.extractAndExpand(nh.salt, nil); err != nil {
|
|
return nil, fmt.Errorf("failed to extract final keys: %w", err)
|
|
} else if writeKey, err := newCipher(write); err != nil {
|
|
return nil, fmt.Errorf("failed to create final write cipher: %w", err)
|
|
} else if readKey, err := newCipher(read); err != nil {
|
|
return nil, fmt.Errorf("failed to create final read cipher: %w", err)
|
|
} else if ns, err := newNoiseSocket(fs, writeKey, readKey, frameHandler, disconnectHandler); err != nil {
|
|
return nil, fmt.Errorf("failed to create noise socket: %w", err)
|
|
} else {
|
|
return ns, nil
|
|
}
|
|
}
|
|
|
|
func (nh *NoiseHandshake) MixSharedSecretIntoKey(priv, pub [32]byte) error {
|
|
secret, err := curve25519.X25519(priv[:], pub[:])
|
|
if err != nil {
|
|
return fmt.Errorf("failed to do x25519 scalar multiplication: %w", err)
|
|
}
|
|
return nh.MixIntoKey(secret)
|
|
}
|
|
|
|
func (nh *NoiseHandshake) MixIntoKey(data []byte) error {
|
|
nh.counter = 0
|
|
write, read, err := nh.extractAndExpand(nh.salt, data)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to extract keys for mixing: %w", err)
|
|
}
|
|
nh.salt = write
|
|
nh.key, err = newCipher(read)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create new cipher while mixing keys: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (nh *NoiseHandshake) extractAndExpand(salt, data []byte) (write []byte, read []byte, err error) {
|
|
h := hkdf.New(sha256.New, data, salt, nil)
|
|
write = make([]byte, 32)
|
|
read = make([]byte, 32)
|
|
|
|
if _, err = io.ReadFull(h, write); err != nil {
|
|
err = fmt.Errorf("failed to read write key: %w", err)
|
|
} else if _, err = io.ReadFull(h, read); err != nil {
|
|
err = fmt.Errorf("failed to read read key: %w", err)
|
|
}
|
|
return
|
|
}
|