github.com/coreos/goproxy@v0.0.0-20190513173959-f8dc2d7ba04e/counterecryptor.go (about)

     1  package goproxy
     2  
     3  import (
     4  	"crypto/aes"
     5  	"crypto/cipher"
     6  	"crypto/rsa"
     7  	"crypto/sha256"
     8  	"crypto/x509"
     9  	"errors"
    10  )
    11  
    12  type CounterEncryptorRand struct {
    13  	cipher  cipher.Block
    14  	counter []byte
    15  	rand    []byte
    16  	ix      int
    17  }
    18  
    19  func NewCounterEncryptorRandFromKey(key interface{}, seed []byte) (r CounterEncryptorRand, err error) {
    20  	var keyBytes []byte
    21  	switch key := key.(type) {
    22  	case *rsa.PrivateKey:
    23  		keyBytes = x509.MarshalPKCS1PrivateKey(key)
    24  	default:
    25  		err = errors.New("only RSA keys supported")
    26  		return
    27  	}
    28  	h := sha256.New()
    29  	if r.cipher, err = aes.NewCipher(h.Sum(keyBytes)[:aes.BlockSize]); err != nil {
    30  		return
    31  	}
    32  	r.counter = make([]byte, r.cipher.BlockSize())
    33  	if seed != nil {
    34  		copy(r.counter, h.Sum(seed)[:r.cipher.BlockSize()])
    35  	}
    36  	r.rand = make([]byte, r.cipher.BlockSize())
    37  	r.ix = len(r.rand)
    38  	return
    39  }
    40  
    41  func (c *CounterEncryptorRand) Seed(b []byte) {
    42  	if len(b) != len(c.counter) {
    43  		panic("SetCounter: wrong counter size")
    44  	}
    45  	copy(c.counter, b)
    46  }
    47  
    48  func (c *CounterEncryptorRand) refill() {
    49  	c.cipher.Encrypt(c.rand, c.counter)
    50  	for i := 0; i < len(c.counter); i++ {
    51  		if c.counter[i]++; c.counter[i] != 0 {
    52  			break
    53  		}
    54  	}
    55  	c.ix = 0
    56  }
    57  
    58  func (c *CounterEncryptorRand) Read(b []byte) (n int, err error) {
    59  	if c.ix == len(c.rand) {
    60  		c.refill()
    61  	}
    62  	if n = len(c.rand) - c.ix; n > len(b) {
    63  		n = len(b)
    64  	}
    65  	copy(b, c.rand[c.ix:c.ix+n])
    66  	c.ix += n
    67  	return
    68  }