encryption and decryption of file

Author: Soumo Sarkar

crypto.go

@@ -0,0 +1,100 @@
+package main
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"io"
+)
+
+func newEncryptionKey() []byte {
+	keyBuf := make([]byte, 32)
+	io.ReadFull(rand.Reader, keyBuf)
+	return keyBuf
+}
+
+func copyDecrypt(key []byte, src io.Reader, dst io.Writer) (int, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return 0, err
+	}
+
+	/*
+	   Read the IV from the given io.Reader
+	   in our case which should be the block.BlockSize bytes we read.
+	*/
+	iv := make([]byte, block.BlockSize())
+	if _, err := src.Read(iv); err != nil {
+		return 0, err
+	}
+
+	var (
+		buf    = make([]byte, 32*1024)
+		stream = cipher.NewCTR(block, iv)
+		nw     = block.BlockSize()
+	)
+
+	for {
+		n, err := src.Read(buf)
+		if n > 0 {
+			stream.XORKeyStream(buf, buf[:n])
+			n, err := dst.Write(buf[:n])
+			if err != nil {
+				return 0, err
+			}
+			nw += n
+		}
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	return nw, nil
+}
+
+func copyEncrypt(key []byte, src io.Reader, dst io.Writer) (int, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return 0, err
+	}
+
+	iv := make([]byte, block.BlockSize()) // 16 bytes
+	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
+		return 0, err
+	}
+
+	// prepend the IV to the file
+	if _, err := dst.Write(iv); err != nil {
+		return 0, err
+	}
+
+	var (
+		buf    = make([]byte, 32*1024)
+		stream = cipher.NewCTR(block, iv)
+		nw     = block.BlockSize()
+	)
+
+	for {
+		n, err := src.Read(buf)
+		if n > 0 {
+			stream.XORKeyStream(buf, buf[:n])
+			n, err := dst.Write(buf[:n])
+			if err != nil {
+				return 0, err
+			}
+			nw += n
+		}
+
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	return nw, nil
+}

crypto_test.go

@@ -0,0 +1,35 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+)
+
+func TestCopyEncryptDecrypt(t *testing.T) {
+	payload := "Foo not bar"
+	src := bytes.NewReader([]byte(payload))
+	dst := new(bytes.Buffer)
+	key := newEncryptionKey()
+	_, err := copyEncrypt(key, src, dst)
+	if err != nil {
+		t.Error(err)
+	}
+
+	fmt.Println(len(payload))
+	fmt.Println(len(dst.String()))
+
+	out := new(bytes.Buffer)
+	nw, err := copyDecrypt(key, dst, out)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if nw != 16+len(payload) {
+		t.Fail()
+	}
+
+	if out.String() != payload {
+		t.Errorf("decryption failed!")
+	}
+}

go.mod

@@ -2,9 +2,10 @@ module github.com/SuperSection/FileStorage
 
 go 1.22.5
 
+require github.com/stretchr/testify v1.9.0
+
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/stretchr/testify v1.9.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -4,6 +4,7 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

main.go

@@ -1,8 +1,7 @@
 package main
 
 import (
-	"fmt"
-	"io"
+	"bytes"
 	"log"
 	"time"
 
@@ -20,6 +19,7 @@ func makeServer(listenAddr string, nodes ...string) *FileServer {
 	tcpTransport := p2p.NewTCPTransport(tcpOptions)
 
 	fileServerOptions := FileServerOptions{
+		EncKey:            newEncryptionKey(),
 		StorageRoot:       listenAddr + "_network",
 		PathTransformFunc: CASPathTransformFunc,
 		Transport:         tcpTransport,
@@ -47,19 +47,18 @@ func main() {
 	go s2.Start()
 	time.Sleep(2 * time.Second)
 
-	// data := bytes.NewReader([]byte("my big data file here!"))
-	// s2.Store("cool_picture.jpg", data)
-	// time.Sleep(time.Millisecond * 500)
-
-	r, err := s2.Get("cool_picture.jpg")
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	b, err := io.ReadAll(r)
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	fmt.Println(string(b))
+	data := bytes.NewReader([]byte("my big data file here!"))
+	s2.Store("cool_picture.jpg", data)
+
+	// r, err := s2.Get("cool_picture.jpg")
+	// if err != nil {
+	// 	log.Fatal(err)
+	// }
+	//
+	// b, err := io.ReadAll(r)
+	// if err != nil {
+	// 	log.Fatal(err)
+	// }
+	//
+	// fmt.Println(string(b))
 }

p2p/encoding.go

@@ -41,4 +41,3 @@ func (dec *DefaultDecoder) Decode(r io.Reader, msg *RPC) error {
 
 	return nil
 }
-

p2p/tcp_transport_test.go

@@ -7,17 +7,13 @@ import (
 )
 
 func TestTCPTransport(t *testing.T) {
-	tcpOptions := TCPTransportOptions{
-		ListenAddress: ":4001",
-		Decoder:       &DefaultDecoder{},
+	opts := TCPTransportOptions{
+		ListenAddress: ":3000",
 		HandshakeFunc: NOPHandshakeFunc,
+		Decoder:       &DefaultDecoder{},
 	}
+	tr := NewTCPTransport(opts)
+	assert.Equal(t, tr.ListenAddr, ":3000")
 
-	tr := NewTCPTransport(tcpOptions)
-	assert.Equal(t, tr.ListenAddress, ":4001")
-
-	// Server
 	assert.Nil(t, tr.ListenAndAccept())
-
-	select {}
-}
+}

server.go

@@ -14,6 +14,7 @@ import (
 )
 
 type FileServerOptions struct {
+	EncKey            []byte
 	StorageRoot       string
 	PathTransformFunc PathTransformFunc
 	Transport         p2p.Transport
@@ -144,7 +145,7 @@ func (server *FileServer) Store(key string, r io.Reader) error {
 	msg := Message{
 		Payload: MessageStoreFile{
 			Key:  key,
-			Size: size,
+			Size: size + 16,
 		},
 	}
 
@@ -157,7 +158,7 @@ func (server *FileServer) Store(key string, r io.Reader) error {
 	// TODO: use a multiwriter here.
 	for _, peer := range server.peers {
 		peer.Send([]byte{p2p.IncomingStream})
-		n, err := io.Copy(peer, fileBuffer)
+		n, err := copyEncrypt(server.EncKey, fileBuffer, peer)
 		if err != nil {
 			return err
 		}