github.com/AntonOrnatskyi/goproxy@v0.0.0-20190205095733-4526a9fa18b4/utils/conncrypt/conncrypt.go (about)

     1  package conncrypt
     2  
     3  import (
     4  	"crypto/aes"
     5  	"crypto/cipher"
     6  	"crypto/sha256"
     7  	"hash"
     8  	"io"
     9  	"net"
    10  
    11  	"golang.org/x/crypto/pbkdf2"
    12  )
    13  
    14  //Confg defaults
    15  const DefaultIterations = 2048
    16  const DefaultKeySize = 32 //256bits
    17  var DefaultHashFunc = sha256.New
    18  var DefaultSalt = []byte(`
    19  (;QUHj.BQ?RXzYSO]ifkXp/G!kFmWyXyEV6Nt!d|@bo+N$L9+<d$|g6e26T}
    20  Ao<:>SOd,6acYKY_ec+(x"R";\'4&fTAVu92GVA-wxBptOTM^2,iP5%)wnhW
    21  hwk=]Snsgymt!3gbP2pe=J//}1a?lp9ej=&TB!C_V(cT2?z8wyoL_-13fd[]
    22  `) //salt must be predefined in order to derive the same key
    23  
    24  //Config stores the PBKDF2 key generation parameters
    25  type Config struct {
    26  	Password   string
    27  	Salt       []byte
    28  	Iterations int
    29  	KeySize    int
    30  	HashFunc   func() hash.Hash
    31  }
    32  
    33  //New creates an AES encrypted net.Conn by generating
    34  //a key using PBKDF2 with the provided configuration
    35  func New(conn net.Conn, c *Config) net.Conn {
    36  	//set defaults
    37  	if len(c.Salt) == 0 {
    38  		c.Salt = DefaultSalt
    39  	}
    40  	if c.Iterations == 0 {
    41  		c.Iterations = DefaultIterations
    42  	}
    43  	if c.KeySize != 16 && c.KeySize != 24 && c.KeySize != 32 {
    44  		c.KeySize = DefaultKeySize
    45  	}
    46  	if c.HashFunc == nil {
    47  		c.HashFunc = DefaultHashFunc
    48  	}
    49  
    50  	//generate key
    51  	key := pbkdf2.Key([]byte(c.Password), c.Salt, c.Iterations, c.KeySize, c.HashFunc)
    52  
    53  	// could use scrypt, but it's a bit slow...
    54  	// dk, err := scrypt.Key([]byte(c.Password), c.Salt, 16384, 8, 1, 32)
    55  
    56  	//key will be always be the correct size so this will never error
    57  	conn, _ = NewFromKey(conn, key)
    58  	return conn
    59  }
    60  
    61  //NewFromKey creates an AES encrypted net.Conn using the provided key
    62  func NewFromKey(conn net.Conn, key []byte) (net.Conn, error) {
    63  	block, err := aes.NewCipher([]byte(key))
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  	//hash(key) -> read IV
    68  	riv := DefaultHashFunc().Sum(key)
    69  	rstream := cipher.NewCFBDecrypter(block, riv[:aes.BlockSize])
    70  	reader := &cipher.StreamReader{S: rstream, R: conn}
    71  	//hash(read IV) -> write IV
    72  	wiv := DefaultHashFunc().Sum(riv)
    73  	wstream := cipher.NewCFBEncrypter(block, wiv[:aes.BlockSize])
    74  	writer := &cipher.StreamWriter{S: wstream, W: conn}
    75  
    76  	return &cryptoConn{
    77  		Conn: conn,
    78  		r:    reader,
    79  		w:    writer,
    80  	}, nil
    81  }
    82  
    83  type cryptoConn struct {
    84  	net.Conn
    85  	r io.Reader
    86  	w io.Writer
    87  }
    88  
    89  //replace read and write methods
    90  func (c *cryptoConn) Read(p []byte) (int, error) {
    91  	return c.r.Read(p)
    92  }
    93  func (c *cryptoConn) Write(p []byte) (int, error) {
    94  	return c.w.Write(p)
    95  }