Add hashing of local file path to translate while distributing across other servers

Author: Soumo Sarkar

crypto.go

@@ -3,35 +3,33 @@ package main
 import (
 	"crypto/aes"
 	"crypto/cipher"
+	"crypto/md5"
 	"crypto/rand"
+	"encoding/hex"
 	"io"
 )
 
+func generateID() string {
+	buf := make([]byte, 32)
+	io.ReadFull(rand.Reader, buf)
+	return hex.EncodeToString(buf)
+}
+
+func hashKey(key string) string {
+	hash := md5.Sum([]byte(key))
+	return hex.EncodeToString(hash[:])
+}
+
 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
-	}
-
+func copyStream(stream cipher.Stream, BlockSize int, src io.Reader, dst io.Writer) (int, error) {
 	var (
-		buf    = make([]byte, 32*1024)
-		stream = cipher.NewCTR(block, iv)
-		nw     = block.BlockSize()
+		buf = make([]byte, 32*1024)
+		nw  = BlockSize
 	)
 
 	for {
@@ -55,6 +53,25 @@ func copyDecrypt(key []byte, src io.Reader, dst io.Writer) (int, error) {
 	return nw, nil
 }
 
+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
+	}
+
+	stream := cipher.NewCTR(block, iv)
+	return copyStream(stream, block.BlockSize(), src, dst)
+}
+
 func copyEncrypt(key []byte, src io.Reader, dst io.Writer) (int, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
@@ -71,30 +88,6 @@ func copyEncrypt(key []byte, src io.Reader, dst io.Writer) (int, error) {
 		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
+	stream := cipher.NewCTR(block, iv)
+	return copyStream(stream, block.BlockSize(), src, dst)
 }

main.go

@@ -2,6 +2,8 @@ package main
 
 import (
 	"bytes"
+	"fmt"
+	"io"
 	"log"
 	"time"
 
@@ -36,29 +38,37 @@ func makeServer(listenAddr string, nodes ...string) *FileServer {
 func main() {
 
 	s1 := makeServer(":3000", "")
-	s2 := makeServer(":4000", ":3000")
+	s2 := makeServer(":7000", ":3000")
+	s3 := makeServer(":5000", ":3000", ":7000")
 
-	go func() {
-		log.Fatal(s1.Start())
-	}()
+	go func() { log.Fatal(s1.Start()) }()
+	time.Sleep(500 * time.Millisecond)
+	go func() { log.Fatal(s2.Start()) }()
 
 	time.Sleep(2 * time.Second)
 
-	go s2.Start()
+	go s3.Start()
 	time.Sleep(2 * time.Second)
 
-	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))
+	for i := 0; i < 20; i++ {
+		key := fmt.Sprintf("picture_%d.png", i)
+		data := bytes.NewReader([]byte("my big data file here!"))
+		s3.Store(key, data)
+
+		if err := s3.storage.Delete(s3.ID, key); err != nil {
+			log.Fatal(err)
+		}
+
+		r, err := s3.Get(key)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		b, err := io.ReadAll(r)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		fmt.Println(string(b))
+	}
 }

p2p/tcp_transport_test.go

@@ -13,7 +13,7 @@ func TestTCPTransport(t *testing.T) {
 		Decoder:       &DefaultDecoder{},
 	}
 	tr := NewTCPTransport(opts)
-	assert.Equal(t, tr.ListenAddr, ":3000")
+	assert.Equal(t, tr.ListenAddress, ":3000")
 
 	assert.Nil(t, tr.ListenAndAccept())
 }

server.go

@@ -14,6 +14,7 @@ import (
 )
 
 type FileServerOptions struct {
+	ID                string
 	EncKey            []byte
 	StorageRoot       string
 	PathTransformFunc PathTransformFunc
@@ -37,6 +38,10 @@ func NewFileServer(opts FileServerOptions) *FileServer {
 		PathTransformFunc: opts.PathTransformFunc,
 	}
 
+	if len(opts.ID) == 0 {
+		opts.ID = generateID()
+	}
+
 	return &FileServer{
 		FileServerOptions: opts,
 		storage:           NewStorage(storageOpts),
@@ -45,16 +50,6 @@ func NewFileServer(opts FileServerOptions) *FileServer {
 	}
 }
 
-func (server *FileServer) stream(msg *Message) error {
-	peers := []io.Writer{}
-	for _, peer := range server.peers {
-		peers = append(peers, peer)
-	}
-
-	mw := io.MultiWriter(peers...)
-	return gob.NewEncoder(mw).Encode(msg)
-}
-
 func (server *FileServer) broadcast(msg *Message) error {
 	/* encode the message */
 	buf := new(bytes.Buffer)
@@ -78,26 +73,29 @@ type Message struct {
 }
 
 type MessageStoreFile struct {
+	ID   string
 	Key  string
 	Size int64
 }
 
 type MessageGetFile struct {
+	ID  string
 	Key string
 }
 
 func (server *FileServer) Get(key string) (io.Reader, error) {
-	if server.storage.Has(key) {
+	if server.storage.Has(server.ID, key) {
 		fmt.Printf("[%s] serving file (%s) from local disk\n", server.Transport.ListenAddr(), key)
-		_, r, err := server.storage.Read(key)
+		_, r, err := server.storage.Read(server.ID, key)
 		return r, err
 	}
 
 	fmt.Printf("[%s] don't have file (%s) locally, fetching from network...\n", server.Transport.ListenAddr(), key)
 
 	msg := Message{
 		Payload: MessageGetFile{
-			Key: key,
+			ID:  server.ID,
+			Key: hashKey(key),
 		},
 	}
 
@@ -109,10 +107,11 @@ func (server *FileServer) Get(key string) (io.Reader, error) {
 
 	for _, peer := range server.peers {
 		/* first read the file size so we can limit the amount of bytes
-		   that wee need from the connection so it will not keep hanging */
+		   that we need from the connection so it will not keep hanging */
 		var fileSize int64
 		binary.Read(peer, binary.LittleEndian, &fileSize)
-		n, err := server.storage.Write(key, io.LimitReader(peer, 22))
+
+		n, err := server.storage.writeDecrypt(server.EncKey, server.ID, key, io.LimitReader(peer, fileSize))
 		if err != nil {
 			return nil, err
 		}
@@ -122,7 +121,7 @@ func (server *FileServer) Get(key string) (io.Reader, error) {
 		peer.CloseStream()
 	}
 
-	_, r, err := server.storage.Read(key)
+	_, r, err := server.storage.Read(server.ID, key)
 	return r, err
 }
 
@@ -137,14 +136,15 @@ func (server *FileServer) Store(key string, r io.Reader) error {
 		tee        = io.TeeReader(r, fileBuffer)
 	)
 
-	size, err := server.storage.Write(key, tee)
+	size, err := server.storage.Write(server.ID, key, tee)
 	if err != nil {
 		return nil
 	}
 
 	msg := Message{
 		Payload: MessageStoreFile{
-			Key:  key,
+			ID:   server.ID,
+			Key:  hashKey(key),
 			Size: size + 16,
 		},
 	}
@@ -155,17 +155,20 @@ func (server *FileServer) Store(key string, r io.Reader) error {
 
 	time.Sleep(time.Millisecond * 5)
 
-	// TODO: use a multiwriter here.
+	peers := []io.Writer{}
 	for _, peer := range server.peers {
-		peer.Send([]byte{p2p.IncomingStream})
-		n, err := copyEncrypt(server.EncKey, fileBuffer, peer)
-		if err != nil {
-			return err
-		}
+		peers = append(peers, peer)
+	}
 
-		fmt.Println("received and written to disk: ", n)
+	mw := io.MultiWriter(peers...)
+	mw.Write([]byte{p2p.IncomingStream})
+	n, err := copyEncrypt(server.EncKey, fileBuffer, mw)
+	if err != nil {
+		return err
 	}
 
+	fmt.Printf("[%s] received and written (%d) bytes to disk\n", server.Transport.ListenAddr(), n)
+
 	return nil
 }
 
@@ -219,13 +222,13 @@ func (server *FileServer) handleMessage(from string, msg *Message) error {
 }
 
 func (server *FileServer) handleMessageGetFile(from string, msg MessageGetFile) error {
-	if !server.storage.Has(msg.Key) {
+	if !server.storage.Has(msg.ID, msg.Key) {
 		return fmt.Errorf("[%s] need to serve file (%s) but it doesn't exist on disk", server.Transport.ListenAddr(), msg.Key)
 	}
 
 	fmt.Printf("[%s] serving file (%s) over the network\n", server.Transport.ListenAddr(), msg.Key)
 
-	fileSize, r, err := server.storage.Read(msg.Key)
+	fileSize, r, err := server.storage.Read(msg.ID, msg.Key)
 	if err != nil {
 		return err
 	}
@@ -261,7 +264,7 @@ func (server *FileServer) handleMessageStoreFile(from string, msg MessageStoreFi
 		return fmt.Errorf("peer (%s) could not be found in the peer list", from)
 	}
 
-	n, err := server.storage.Write(msg.Key, io.LimitReader(peer, msg.Size))
+	n, err := server.storage.Write(msg.ID, msg.Key, io.LimitReader(peer, msg.Size))
 	if err != nil {
 		return err
 	}
@@ -280,7 +283,7 @@ func (server *FileServer) bootstrapNetwork() error {
 		}
 
 		go func(addr string) {
-			fmt.Println("attempting to connect with remote: ", addr)
+			fmt.Printf("[%s] attempting to connect with remote %s\n", server.Transport.ListenAddr(), addr)
 
 			if err := server.Transport.Dial(addr); err != nil {
 				log.Println("dial error: ", err)
@@ -292,6 +295,7 @@ func (server *FileServer) bootstrapNetwork() error {
 }
 
 func (server *FileServer) Start() error {
+	fmt.Printf("[%s] starting fileserver...\n", server.Transport.ListenAddr())
 	if err := server.Transport.ListenAndAccept(); err != nil {
 		return err
 	}