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 } 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, }, }) }