codename-sendaround/uploader.go

package sendaround

import (
	"io"
	"net/url"
	"sync"

	"github.com/icedream/sendaround/internal"
)

// NOTE - Any larger value than this seems to make data transmission fail within the pion/webrtc library.
// NOTE - Max WebRTC packet size is 0xffff in any case.
const DefaultCopyBufferSize = 16 * 1024

type Uploader struct {
	dataChannel *internal.DataChannel

	filesLock sync.RWMutex
	files     map[string]File

	err error

	stateLock sync.RWMutex
	state     ConnectionState
	stateC    chan ConnectionState

	token string
	url   *url.URL

	CopyBufferSize int
}

func (server *Uploader) init() {
	if server.CopyBufferSize == 0 {
		server.CopyBufferSize = DefaultCopyBufferSize
	}
	server.stateC = make(chan ConnectionState, 1)
	server.files = map[string]File{}
	server.changeState(ConnectionState{Type: WaitingForClient})

	server.dataChannel.OnOpen(func() {
		server.filesLock.RLock()
		defer server.filesLock.RUnlock()
		server.changeState(ConnectionState{Type: Connected})

		for _, f := range server.files {
			server.dataChannel.SendMessage(&internal.Message{
				FileOfferMessage: &internal.FileOfferMessage{
					FileName: f.FileName(),
					MimeType: f.MimeType(),
					Length:   f.Length(),
				},
			})
		}

		server.dataChannel.SendMessage(&internal.Message{
			SessionInitializedMessage: &internal.SessionInitializedMessage{},
		})
	})

	server.dataChannel.OnClose(func() {
		switch server.state.Type {
		case TransmittingFile:
			server.changeState(ConnectionState{Type: Failed, Error: ErrClientClosedConnection})
		case Failed:
		case Disconnected:
		default:
			server.changeState(ConnectionState{Type: Disconnected})
		}
		close(server.stateC)
	})

	server.dataChannel.OnError(func(err error) {
		server.abortConnection(err)
	})

	server.dataChannel.OnSendaroundMessage(func(msg *internal.Message) {
		switch {
		case msg.FileTransferRequestMessage != nil:
			if err := server.sendFile(msg.FileTransferRequestMessage.FileName); err != nil {
				server.abortConnection(err)
			}
		default:
			// Something's wrong with this message...
		}
	})
}

func (server *Uploader) StateC() chan ConnectionState {
	return server.stateC
}

func (server *Uploader) Token() string {
	return server.token
}

func (server *Uploader) URL() *url.URL {
	return server.url
}

func (server *Uploader) Close() error {
	server.stateLock.Lock()
	defer server.stateLock.Unlock()

	if server.state.Type == Failed || server.state.Type == Disconnected {
		return nil
	}

	if server.state.Type == TransmittingFile {
		return ErrInvalidState
	}

	server.dataChannel.Close()
	return nil
}

func (server *Uploader) abortConnection(err error) {
	server.stateLock.Lock()
	defer server.stateLock.Unlock()

	state := server.state
	state.Type = Failed
	state.Error = err
	server.changeState(state)

	server.dataChannel.Close()
}

func (server *Uploader) _unlocked_changeState(state ConnectionState) {
	server.state = state
	server.stateC <- state
}

func (server *Uploader) changeState(state ConnectionState) {
	server.stateLock.Lock()
	defer server.stateLock.Unlock()
	server._unlocked_changeState(state)
}

func (server *Uploader) sendFile(path string) (err error) {
	server.stateLock.Lock()
	defer server.stateLock.Unlock()

	if server.state.Type != Connected {
		err = ErrInvalidState
		return
	}

	f, ok := server.files[normalizeFilePath(path)]
	if !ok {
		err = ErrFileNotFound
		return
	}

	state := ConnectionState{Type: TransmittingFile, CurrentFile: f}
	defer func() {
		if err != nil {
			server.abortConnection(err)
		} else {
			state.Type = Connected
			server._unlocked_changeState(state)
		}
	}()

	server._unlocked_changeState(state)
	r := f.getReader()

	var n int
	b := make([]byte, server.CopyBufferSize)
	for {
		n, err = r.Read(b)
		if err == io.EOF {
			err = nil
			break
		}

		n, err = server.dataChannel.Write(b[0:n])
		if err != nil {
			return
		}

		state.TransmittedLength += uint64(n)
		server._unlocked_changeState(state)
	}
	return
}

func (server *Uploader) AddFile(fileToAdd File) (err error) {
	server.filesLock.Lock()
	defer server.filesLock.Unlock()

	if _, exists := server.files[fileToAdd.FileName()]; exists {
		err = ErrFileAlreadyExists
		return
	}

	server.files[fileToAdd.FileName()] = fileToAdd

	// If already connected, sync this new entry immediately
	server.stateLock.RLock()
	defer server.stateLock.RUnlock()
	if server.state.Type != WaitingForClient && server.state.Type != Failed && server.state.Type != Disconnected {
		server.dataChannel.SendMessage(&internal.Message{
			FileOfferMessage: &internal.FileOfferMessage{
				FileName: fileToAdd.FileName(),
				MimeType: fileToAdd.MimeType(),
				Length:   fileToAdd.Length(),
			},
		})
	}

	return
}

func (server *Uploader) RemoveFile(filePath string) {
	server.filesLock.Lock()
	defer server.filesLock.Unlock()

	delete(server.files, normalizeFilePath(filePath))

	// If already connected, sync this removed entry immediately
	server.stateLock.RLock()
	defer server.stateLock.RUnlock()
	server.dataChannel.SendMessage(&internal.Message{
		FileUnofferMessage: &internal.FileUnofferMessage{
			FileName: filePath,
		},
	})
}