1 files changed, 13 insertions, 9 deletions
diff --git a/PumpKINDlg.cpp b/PumpKINDlg.cpp
index 070f3e8..2a01918 100644
--- a/PumpKINDlg.cpp
+++ b/PumpKINDlg.cpp
@@ -534,100 +534,100 @@ UINT rv = 0;
        return rv;
 }
 tftp::tftp()
 {
        length=0;
 }
 void CXferSocket::OnSend(int nErrorCode)
 {
        if(nErrorCode){
                ASSERT(m_Daddy);
                m_Daddy->LogLine(IDS_LOG_XFERSEND);
                return;
        }
        if(!m_Queue.IsEmpty()){
        tftp *p = m_Queue.GetHead();
                ASSERT(p);
                m_Queue.RemoveHead();
                if(!p->Send(this,&m_Peer)){
                        ASSERT(m_Daddy);
                        m_Daddy->LogLine(IDS_LOG_XFERUDPSEND);
                }
                delete p;
        }
        DoSelect();
        if(m_Queue.IsEmpty()){
                switch(state){
                case stateDeny:
                        Destroy(FALSE);
                        break;
                case stateFinish:
                        Destroy(TRUE);
                        break;
                }
        }
 }
 BOOL tftp::Send(CAsyncSocket *socket,SOCKADDR_IN* saddr)
 {
        ASSERT(socket);
 int rv = socket->SendTo(udpBase(),length,(SOCKADDR*)saddr,sizeof(SOCKADDR_IN));
        if(rv!=length)
                return FALSE;
        return TRUE;
 }
-void CXferSocket::DoSelect()
+void CXferSocket::DoSelect(BOOL do_select)
 {
        if(m_Peer.sin_addr.s_addr!=INADDR_NONE)
-                AsyncSelect(FD_CLOSE|FD_READ|(m_Queue.IsEmpty()?0:FD_WRITE));
+                AsyncSelect(FD_CLOSE|FD_READ|((m_Queue.IsEmpty()&&!do_select)?0:FD_WRITE));
 }
 void CXferSocket::OnReceive(int nErrorCode)
 {
        if(nErrorCode){
                ASSERT(m_Daddy);
                m_Daddy->LogLine(IDS_LOG_XFERRECEIVE);
                return;
        }
        ASSERT(m_Daddy);
 DWORD fionread = 0;
        VERIFY(IOCtl(FIONREAD,&fionread));
 tftp *p = tftp::Allocate(fionread);
        ASSERT(p);
 SOCKADDR_IN sin;
        if(!p->Receive(this,fionread,&sin)){
                m_Daddy->LogLine(IDS_LOG_XFERUDPRECEIVE);
                delete p;
        }else
                if(m_Peer.sin_addr.s_addr==INADDR_NONE){
                        m_Peer.sin_addr=sin.sin_addr;
                        m_Peer.sin_port=sin.sin_port;
        }
 BOOL alive = TRUE;
        if(state==stateInit){
                state=stateXfer;
                m_Peer.sin_port=sin.sin_port;
                UpdateList();
        }
        if(sin.sin_addr.s_addr!=m_Peer.sin_addr.s_addr || sin.sin_port!=m_Peer.sin_port){
                m_Daddy->LogLine(IDS_LOG_XFERSOURCETID);
                // *** Bounce it!
        }else{
                alive = OnTFTP(p);
        }
        delete p;
        if(alive){
                DoSelect();
                ResetTimeout();
        }
 }
 void CXferSocket::SetPeer(SOCKADDR_IN *sin)
 {
        ASSERT(sin);
        memmove(&m_Peer,sin,sizeof(m_Peer));
 }
@@ -809,211 +809,215 @@ tftp* tftp::Allocate(UINT tftpSize)
        ASSERT(tftpSize);
 tftp* rv = (tftp*) new BYTE[tftpSlackSize+tftpSize];
        ASSERT(rv);
        rv->length=tftpSize;
        return rv;
 }
 void tftp::errSet(UINT code,LPCTSTR msg)
 {
        ASSERT(this);
        ASSERT(length>=data.m_ERROR.tftpSize(msg));
        strcpy((char*)data.m_ERROR.data,msg);
        data.m_ERROR.SetCode(code);
 }
 void CXferSocket::PostTFTP(tftp* p,BOOL retryable)
 {
        ASSERT(p);
        m_Queue.AddTail(p);
        DoSelect();
        if(!m_bRetry){
                if(retryable)
                        SetTry(p);
                else
                        SetTry();
        }
        ResetTimeout();
 }
 void CXferSocket::Deny(UINT errCode,UINT errID)
 {
        PostError(errCode,errID);
        state=stateDeny;
 }
 void CRRQSocket::DoXfer()
 {
 tftp *p = tftp::Allocate(tftp::tftpDATA::tftpSize(m_blkSize));
        ASSERT(p);
        p->SetOpcode(tftp::opDATA);
        TRY{
                m_File.Seek(m_ACK*m_blkSize,CFile::begin);
        int bytes = m_File.Read(p->data.m_DATA.data,m_blkSize);
                p->data.m_DATA.SetBlock(m_ACK+1);
                p->length=p->length-m_blkSize+bytes;
                m_LastSlack = m_blkSize-bytes;
                PostTFTP(p);
                if(bytes<m_blkSize){
-                        state=stateFinish;
+                        state=stateClosing; m_ACKtoClose = m_ACK+1;
-                        ASSERT(m_Daddy);
-                CString tmp;
-                        tmp.Format(IDS_LOG_XFERRRQFINISHED,(LPCTSTR)m_FileName);
-                        m_Daddy->LogLine(tmp);
                }
        }CATCH(CFileException,e){
                Deny(e);
        }END_CATCH
 }
 UINT tftp::tftpDATA::tftpSize(UINT blkSize)
 {
        return tftpHdrSize-tftpSlackSize+sizeof(tftp::tftpDATA)
                -sizeof(BYTE)+blkSize;
 }
 void CXferSocket::Deny(CFileException* e)
 {
        PostError(e);
        state=stateDeny;
 }
 void CXferSocket::PostError(UINT errCode,UINT errID)
 {
 CString msg;
        msg.LoadString(errID);
        ASSERT(m_Daddy);
        /*// ***
 CString tmp;
        tmp.Format(IDS_LOG_SENTTFTPERROR,errCode,(LPCTSTR)msg);
        m_Daddy->LogLine(tmp);
 */
 tftp* err = tftp::Allocate(tftp::tftpERROR::tftpSize(msg));
 err->SetOpcode(tftp::opERROR);
        err->errSet(errCode,msg);
        PostTFTP(err);
 }
 void CXferSocket::PostError(CFileException* e)
 {
 UINT eCode;
 UINT eMsgID;
        switch(e->m_cause){
        case CFileException::fileNotFound:
                eCode=tftp::errNotFound;
                eMsgID=IDS_TFTP_ERROR_NOTFOUND;
                break;
        case CFileException::accessDenied:
                eCode=tftp::errAccessViolation;
                eMsgID=IDS_TFTP_ERROR_ACCESS;
                break;
        case CFileException::directoryFull:
                eCode=tftp::errDiskFull;
                eMsgID=IDS_TFTP_ERROR_DIRFULL;
                break;
        case CFileException::sharingViolation:
                eCode=tftp::errAccessViolation;
                eMsgID=IDS_TFTP_ERROR_SHARING;
                break;
        case CFileException::diskFull:
                eCode=tftp::errDiskFull;
                eMsgID=IDS_TFTP_ERROR_DISKFULL;
                break;
        default:
                eCode=tftp::errUndefined;
                eMsgID=IDS_TFTP_ERROR_UNDEFINED;
                break;
        }
        PostError(eCode,eMsgID);
 }
 ULONG CRRQSocket::GetACK(void)
 {
        return (m_ACK*m_blkSize)-m_LastSlack;
 }
 BOOL CRRQSocket::OnTFTP(tftp* p)
 {
 BOOL rv = TRUE;
        switch(p->Opcode()){
        case tftp::opOACK:
+                {
                m_ACK=0;
                ASSERT(state!=stateFinish);
-        {
        tftp::tftpOptions o;
                if(p->GetOptions(&o)){
                CString v;
                        if(o.Lookup(tftpoBSize,v)){
                                m_blkSize=atoi(v);
                                        if(!m_blkSize){// *** More sanity checks
                                        Deny(tftp::errOption,IDS_TFTP_ERROR_BSIZE);
                                        rv = TRUE;
                                        break;
                                }
                        }
                        if(o.Lookup(tftpoTOut,v)){
                                m_timeOut=atoi(v);
                                        if(!m_timeOut){// *** More sanity checks
                                        Deny(tftp::errOption,IDS_TFTP_ERROR_TOUT);
                                        rv = TRUE;
                                        break;
                                }
                        }
                        if(o.Lookup(tftpoXResume,v)){
                                m_ACK=atoi(v);
                        }
                }
                UpdateList();
                DoXfer();
        }
                break;
        case tftp::opACK:
                m_ACK=p->data.m_ACK.Block();
-                if(state!=stateFinish){
+                if(state==stateClosing && m_ACK==m_ACKtoClose) {
+                        state = stateFinish;
+                        ASSERT(m_Daddy);
+                CString tmp;
+                        tmp.Format(IDS_LOG_XFERRRQFINISHED,(LPCTSTR)m_FileName);
+                        m_Daddy->LogLine(tmp);
+                        rv = FALSE;
+                        DoSelect(TRUE);
+                }else if(state!=stateFinish){
                        UpdateList();
                        DoXfer();
                }
                break;
        case tftp::opERROR:
                {
                        ASSERT(m_Daddy);
                CString tmp;
                        tmp.Format(IDS_LOG_GOTTFTPERROR,p->data.m_ERROR.Code(),(LPCTSTR)p->errMessage());
                        m_Daddy->LogLine(tmp);
                }
                Destroy(FALSE);
                rv = FALSE;
                break;
        default:
                {
                        ASSERT(m_Daddy);
                        CString tmp;
                        tmp.Format(IDS_LOG_XFEROPCODE,p->Opcode());
                        m_Daddy->LogLine(tmp);
                        // *** Self destruct maybe??
                }
                break;
        }
        return rv;
 }
 BOOL CWRQSocket::OnTFTP(tftp* p)
 {
        switch(p->Opcode()){
        case tftp::opOACK:
                ASSERT(state!=stateFinish);
        {
                if(m_bResume)
                        m_ACK=m_File.GetLength()/m_blkSize;
                else
                        m_ACK=0;
        tftp::tftpOptions o;
                if(p->GetOptions(&o)){
                CString v;
                        if(o.Lookup(tftpoBSize,v)){
                                m_blkSize=atoi(v);
                                        if(!m_blkSize){// *** More sanity checks
                                        Deny(tftp::errOption,IDS_TFTP_ERROR_BSIZE);
                                        return TRUE;
                                }
                        }
                        if(o.Lookup(tftpoTOut,v)){