105 lines
1.8 KiB
Go
105 lines
1.8 KiB
Go
package quic
|
|
|
|
import (
|
|
"crypto/aes"
|
|
"crypto/cipher"
|
|
"crypto/rand"
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
|
|
"github.com/lucas-clemente/quic-go"
|
|
)
|
|
|
|
type quicConn struct {
|
|
quic.Session
|
|
quic.Stream
|
|
}
|
|
|
|
func QUICConn(session quic.Session, stream quic.Stream) net.Conn {
|
|
return &quicConn{
|
|
Session: session,
|
|
Stream: stream,
|
|
}
|
|
}
|
|
|
|
type quicCipherConn struct {
|
|
net.PacketConn
|
|
key []byte
|
|
}
|
|
|
|
func QUICCipherConn(conn net.PacketConn, key []byte) net.PacketConn {
|
|
return &quicCipherConn{
|
|
PacketConn: conn,
|
|
key: key,
|
|
}
|
|
}
|
|
|
|
func (conn *quicCipherConn) ReadFrom(data []byte) (n int, addr net.Addr, err error) {
|
|
n, addr, err = conn.PacketConn.ReadFrom(data)
|
|
if err != nil {
|
|
return
|
|
}
|
|
b, err := conn.decrypt(data[:n])
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
copy(data, b)
|
|
|
|
return len(b), addr, nil
|
|
}
|
|
|
|
func (conn *quicCipherConn) WriteTo(data []byte, addr net.Addr) (n int, err error) {
|
|
b, err := conn.encrypt(data)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
_, err = conn.PacketConn.WriteTo(b, addr)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
return len(b), nil
|
|
}
|
|
|
|
func (conn *quicCipherConn) encrypt(data []byte) ([]byte, error) {
|
|
c, err := aes.NewCipher(conn.key)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
gcm, err := cipher.NewGCM(c)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
nonce := make([]byte, gcm.NonceSize())
|
|
if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return gcm.Seal(nonce, nonce, data, nil), nil
|
|
}
|
|
|
|
func (conn *quicCipherConn) decrypt(data []byte) ([]byte, error) {
|
|
c, err := aes.NewCipher(conn.key)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
gcm, err := cipher.NewGCM(c)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
nonceSize := gcm.NonceSize()
|
|
if len(data) < nonceSize {
|
|
return nil, errors.New("ciphertext too short")
|
|
}
|
|
|
|
nonce, ciphertext := data[:nonceSize], data[nonceSize:]
|
|
return gcm.Open(nil, nonce, ciphertext, nil)
|
|
}
|