encode.go 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. // Copyright 2017 fatedier, fatedier@gmail.com
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package crypto
  15. import (
  16. "crypto/aes"
  17. "crypto/cipher"
  18. "crypto/rand"
  19. "crypto/sha1"
  20. "io"
  21. "golang.org/x/crypto/pbkdf2"
  22. )
  23. const (
  24. salt = "frp"
  25. )
  26. // NewWriter returns a new Writer that encrypts bytes to w.
  27. func NewWriter(w io.Writer, key []byte) (*Writer, error) {
  28. key = pbkdf2.Key(key, []byte(salt), 64, aes.BlockSize, sha1.New)
  29. // random iv
  30. iv := make([]byte, aes.BlockSize)
  31. if _, err := io.ReadFull(rand.Reader, iv); err != nil {
  32. return nil, err
  33. }
  34. block, err := aes.NewCipher(key)
  35. if err != nil {
  36. return nil, err
  37. }
  38. return &Writer{
  39. w: w,
  40. enc: &cipher.StreamWriter{
  41. S: cipher.NewCFBEncrypter(block, iv),
  42. W: w,
  43. },
  44. key: key,
  45. iv: iv,
  46. }, nil
  47. }
  48. // Writer is an io.Writer that can write encrypted bytes.
  49. // Now it only support aes-128-cfb.
  50. type Writer struct {
  51. w io.Writer
  52. enc *cipher.StreamWriter
  53. key []byte
  54. iv []byte
  55. ivSend bool
  56. err error
  57. }
  58. // Write satisfies the io.Writer interface.
  59. func (w *Writer) Write(p []byte) (nRet int, errRet error) {
  60. if w.err != nil {
  61. return 0, w.err
  62. }
  63. // When write is first called, iv will be written to w.w
  64. if !w.ivSend {
  65. w.ivSend = true
  66. _, errRet = w.w.Write(w.iv)
  67. if errRet != nil {
  68. w.err = errRet
  69. return
  70. }
  71. }
  72. nRet, errRet = w.enc.Write(p)
  73. if errRet != nil {
  74. w.err = errRet
  75. }
  76. return
  77. }