server.go

package main

import (
	"bytes"
	"encoding/gob"
	"filrestore/p2p"
	"fmt"
	"io"
	"log"
	"math"
	"sync"
	"time"
)

type FileServerConfig struct {
	StorageRoot        string
	PathTransformation PathTransformation
	Transport          p2p.Transport
	BootstrapNodes     []string
}

type FileServer struct {
	Config FileServerConfig

	peersLock sync.Mutex
	peers     map[string]p2p.Peer

	store  *Store
	quitch chan struct{}

	fileWaitCh  map[string]chan Chunk
	waitChMutex sync.Mutex

	fileHMACs map[string][]uint32
	muHMACs   sync.Mutex
}

type FileChunkMessage struct {
	Chunk Chunk
}

type FileGetMessage struct {
	Key string
}

func NewFileServer(config FileServerConfig) *FileServer {
	storeConfig := StoreConfig{
		Root:               config.StorageRoot,
		PathTransformation: config.PathTransformation,
	}
	return &FileServer{
		Config:     config,
		peers:      make(map[string]p2p.Peer),
		store:      NewStore(storeConfig),
		quitch:     make(chan struct{}),
		fileWaitCh: make(map[string]chan Chunk),
		fileHMACs:  make(map[string][]uint32),
	}

}

func (s *FileServer) bootstrapNetwork() {
	for _, addr := range s.Config.BootstrapNodes {
		go func(addr string) {
			if err := s.Config.Transport.Dial(addr); err != nil {
				log.Printf("%s[Peer] %s%s\n", ColorRed, err, ColorReset)
			}
		}(addr)
	}
}

func (s *FileServer) broadcast(p *p2p.RPC) error {

	peers := []io.Writer{}

	for _, peer := range s.peers {
		peers = append(peers, peer)
	}

	log.Println("[Peer] brodcasting rpc to all peers")

	mw := io.MultiWriter(peers...)
	return s.Config.Transport.Encoder().Encode(mw, p)
}

func (s *FileServer) OnPeer(p p2p.Peer) error {
	s.peersLock.Lock()
	defer s.peersLock.Unlock()
	s.peers[p.RemoteAddr().String()] = p

	log.Printf("%s[Peer] remote %s connected%s\n", ColorGreen, p.RemoteAddr(), ColorReset)

	return nil
}

func (s *FileServer) startFileStreamingChannel(key string) chan Chunk {
	s.waitChMutex.Lock()
	ch, found := s.fileWaitCh[key]
	if !found {
		ch = make(chan Chunk)
		s.fileWaitCh[key] = ch
	}
	s.waitChMutex.Unlock()
	return ch
}

func (s *FileServer) streamReceivedChunks(writer *io.PipeWriter, key string, ch chan Chunk) {
	defer writer.Close()
	defer delete(s.fileWaitCh, key)

	var totalchunks int64
	var haveTotal bool
	received := int64(0)

	timeout := 5 * time.Second
	timer := time.NewTimer(timeout)
	defer timer.Stop()

	for !haveTotal || received < totalchunks {

		if !timer.Stop() {
			<-timer.C
		}
		timer.Reset(timeout)

		select {
		case chunk, ok := <-ch:
			if !ok {
				return
			}

			log.Printf("[Chunk] straming (%d) bytes of (%s)\n", chunk.Size, chunk.Key)

			received++
			if !haveTotal {
				totalchunks = int64(math.Ceil(float64(chunk.FileSize) / float64(ChunkSize)))
				haveTotal = true
			}

			if chunk.IsLast {
				writer.Write(chunk.Data)

				log.Printf("%s[Chunk] file (%s) fully chunk-streamed%s\n", ColorGreen, chunk.Key, ColorReset)
				return
			}
			writer.Write(chunk.Data)
		case <-timer.C:
			writer.CloseWithError(fmt.Errorf("%s[Chunk] timeout waiting for the next chunk%s", ColorRed, ColorReset))
			return
		}
	}

}

func (s *FileServer) Get(key string) (io.Reader, error) {
	if _, ok := s.store.Has(key); ok {
		return s.store.Read(key)
	}

	log.Printf("[File] get file of key (%s)\n", key)

	ch := s.startFileStreamingChannel(key)

	p := &p2p.RPC{
		Payload: FileGetMessage{
			Key: key,
		},
	}

	reader, writer := io.Pipe()

	if err := s.broadcast(p); err != nil {
		return nil, err
	}

	go s.streamReceivedChunks(writer, key, ch)

	return reader, nil
}

func (s *FileServer) prepareChunk(key string, buff *bytes.Buffer, filesize int64, chunkindex int, totalchunks int, checksums *[]uint32) (*Chunk, error) {

	sum := chunksum(buff.Bytes())
	*checksums = append(*checksums, sum)

	chunk := &Chunk{
		Key:      key,
		FileSize: filesize,
		Offset:   int64(chunkindex * ChunkSize),
		Size:     int64(buff.Len()),
		Data:     buff.Bytes(),
		Checksum: sum,
		IsLast:   chunkindex+1 == totalchunks,
	}

	if chunk.IsLast {
		chunk.FinalChecksum = finalchecksum(*checksums)
	}

	return chunk, nil
}

func (s *FileServer) Store(key string, r io.Reader, filesize int64) error {
	totalchunks := math.Ceil(float64(filesize) / float64(ChunkSize))

	log.Printf("[File] start storing file of key (%s)\n", key)

	checksums := []uint32{}

	for chunkindex := 0; chunkindex < int(totalchunks); chunkindex++ {
		buff := new(bytes.Buffer)

		_, err := io.CopyN(buff, r, ChunkSize)
		if err != nil && err != io.EOF {
			return err
		}

		chunk, err := s.prepareChunk(key, buff, filesize, chunkindex, int(totalchunks), &checksums)
		if err != nil {
			return err
		}

		log.Printf("[Chunk] writing (%d) bytes chunk of file (%s)\n", chunk.Size, chunk.Key)

		if err := s.store.WriteChunk(*chunk); err != nil {
			return err
		}

		p := &p2p.RPC{
			Payload: FileChunkMessage{
				Chunk: *chunk,
			},
		}
		s.broadcast(p)
	}

	return nil
}

func (s *FileServer) Stop() {
	close(s.quitch)
}

func (s *FileServer) verifyFileChunks(chunk Chunk) error {
	s.muHMACs.Lock()
	defer s.muHMACs.Unlock()

	key := chunk.Key

	sum := chunksum(chunk.Data)
	s.fileHMACs[key] = append(s.fileHMACs[key], sum)

	if sum != chunk.Checksum {
		// ask the peer to resend the corrupted chunk.
		return fmt.Errorf("%s[Chunk] corrupted chunk of file (%s), offset (%d)%s", ColorRed, key, chunk.Offset, ColorReset)
	}

	log.Printf("%s[Chunk] successfully received last=%t chunk %x chunk with no corruption%s\n", ColorGreen, chunk.IsLast, chunk.Checksum, ColorReset)

	if chunk.IsLast {
		finalsum := finalchecksum(s.fileHMACs[key])
		if finalsum != chunk.FinalChecksum {
			// ask the peer to re-stream the whole file
			return fmt.Errorf("%s[File] corrupted file (%s). after checking the finalsum expected (%d), found (%d)%s", ColorRed, key, chunk.FinalChecksum, finalsum, ColorReset)
		}
		log.Printf("%s[File] successfully received the file (%s) no corruption%s\n", ColorGreen, key, ColorReset)

		delete(s.fileHMACs, key)
	}

	return nil
}

func (s *FileServer) handleFileChunkMessage(from string, payload FileChunkMessage) error {
	log.Printf("%s[Chunk] received last=%t chunk %x from %s%s\n", ColorYellow, payload.Chunk.IsLast, payload.Chunk.Checksum, from, ColorReset)

	_, ok := s.peers[from]
	if !ok {
		return fmt.Errorf("%s[Peer] peer %s could not be found in the server peers list%s", ColorRed, from, ColorReset)
	}

	if err := s.verifyFileChunks(payload.Chunk); err != nil {
		return err
	}

	s.waitChMutex.Lock()
	if ch, found := s.fileWaitCh[payload.Chunk.Key]; found {
		ch <- payload.Chunk
	}
	s.waitChMutex.Unlock()

	return s.store.WriteChunk(payload.Chunk)
}

func (s *FileServer) handleFileGetMessage(from string, payload FileGetMessage) error {
	info, ok := s.store.Has(payload.Key)
	if !ok {
		return fmt.Errorf("%s[File] file (%s) couldn't be found on this peer%s", ColorRed, payload.Key, ColorReset)
	}

	log.Printf("%s[File] the file (%s) being searched for from %s is found on peer %s%s\n", ColorYellow, payload.Key, from, s.Config.Transport.ListenAddr(), ColorReset)

	totalchunks := math.Ceil(float64(info.Size()) / float64(ChunkSize))

	peer, ok := s.peers[from]
	if !ok {
		return fmt.Errorf("%s[Peer] peer %s is not found%s", ColorRed, peer, ColorReset)
	}

	checksums := []uint32{}

	for chunkindex := 0; chunkindex < int(totalchunks); chunkindex++ {
		offset := int64(chunkindex * ChunkSize)

		buff, err := s.store.ReadChunk(payload.Key, offset, ChunkSize)
		if err != nil {
			return err
		}

		chunk, err := s.prepareChunk(payload.Key, bytes.NewBuffer(buff), info.Size(), chunkindex, int(totalchunks), &checksums)
		if err != nil {
			return err
		}

		p := &p2p.RPC{
			Payload: FileChunkMessage{
				Chunk: *chunk,
			},
		}

		if err := s.Config.Transport.Encoder().Encode(peer, p); err != nil {
			return err
		}
	}

	return nil
}

func (s *FileServer) handleMessages(message p2p.RPC) error {
	switch p := message.Payload.(type) {
	case FileChunkMessage:
		return s.handleFileChunkMessage(message.From, p)
	case FileGetMessage:
		return s.handleFileGetMessage(message.From, p)
	}
	return nil
}

func (s *FileServer) loop() {
	defer func() {
		log.Printf("%s[File] file server stopped due to error or user quite action%s\n", ColorRed, ColorReset)
		s.Config.Transport.Close()
	}()

	for {
		select {
		case message := <-s.Config.Transport.Consume():
			s.handleMessages(message)
		case <-s.quitch:
			return
		}
	}
}

func (s *FileServer) Start() error {
	if err := s.Config.Transport.ListenAndAccept(); err != nil {
		return err
	}
	s.bootstrapNetwork()
	s.loop()
	return nil
}

func init() {
	gob.Register(FileChunkMessage{})
	gob.Register(FileGetMessage{})
}