Fix an exception caused by file transfer (#56)

CESSProject · Jul 10, 2023 · 6ff0e32 · 6ff0e32
1 parent 187ad41
commit 6ff0e32
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 15 deletions.
diff --git a/core/readfile.go b/core/readfile.go
@@ -57,7 +57,10 @@ func (e *protocols) ReadFileAction(id peer.ID, roothash, datahash, path string,
 	var fstat fs.FileInfo
 	var f *os.File
 	var req pb.ReadfileRequest
-	var resp *pb.ReadfileResponse
+
+	if size != FragmentSize {
+		return errors.New("invalid size")
+	}
 
 	fstat, err = os.Stat(path)
 	if err == nil {
@@ -116,8 +119,15 @@ func (e *protocols) ReadFileAction(id peer.ID, roothash, datahash, path string,
 	// store request so response handler has access to it
 	var respChan = make(chan bool, 1)
 	e.ReadFileProtocol.Lock()
-	e.ReadFileProtocol.requests[req.MessageData.Id] = &readMsgResp{
-		ch: respChan,
+	for {
+		if _, ok := e.ReadFileProtocol.requests[req.MessageData.Id]; ok {
+			req.MessageData.Id = uuid.New().String()
+			continue
+		}
+		e.ReadFileProtocol.requests[req.MessageData.Id] = &readMsgResp{
+			ch: respChan,
+		}
+		break
 	}
 	e.ReadFileProtocol.Unlock()
 	defer func() {
@@ -148,7 +158,22 @@ func (e *protocols) ReadFileAction(id peer.ID, roothash, datahash, path string,
 			return errors.New(ERR_RespTimeOut)
 		}
 
-		resp = e.ReadFileProtocol.requests[req.MessageData.Id].ReadfileResponse
+		e.ReadFileProtocol.Lock()
+		resp, ok := e.ReadFileProtocol.requests[req.MessageData.Id]
+		if !ok {
+			e.ReadFileProtocol.Unlock()
+			return errors.New(ERR_RespFailure)
+		}
+		e.ReadFileProtocol.Unlock()
+
+		if resp.ReadfileResponse == nil {
+			return errors.New(ERR_RespFailure)
+		}
+
+		if len(resp.ReadfileResponse.Data) == 0 || resp.Length == 0 {
+			return errors.New(ERR_RespFailure)
+		}
+
 		num, err = f.Write(resp.Data[:resp.Length])
 		if err != nil {
 			return errors.Wrapf(err, "[write file]")

diff --git a/core/types.go b/core/types.go
@@ -15,8 +15,8 @@ import (
 	ma "github.com/multiformats/go-multiaddr"
 )
 
-const P2PWriteReqRespTime = time.Duration(time.Second * 20)
-const P2PReadReqRespTime = time.Duration(time.Second * 20)
+const P2PWriteReqRespTime = time.Duration(time.Second * 30)
+const P2PReadReqRespTime = time.Duration(time.Second * 30)
 
 const FileProtocolBufSize = 2 * 1024 * 1024
 

diff --git a/core/writefile.go b/core/writefile.go
@@ -58,6 +58,15 @@ func (e *protocols) WriteFileAction(id peer.ID, roothash, path string) error {
 		Roothash:    roothash,
 	}
 
+	fstat, err := os.Stat(path)
+	if err != nil {
+		return err
+	}
+
+	if fstat.Size() != FragmentSize {
+		return errors.New("invalid fragment size")
+	}
+
 	req.Datahash, err = CalcPathSHA256(path)
 	if err != nil {
 		return err
@@ -73,8 +82,15 @@ func (e *protocols) WriteFileAction(id peer.ID, roothash, path string) error {
 	respChan := make(chan bool, 1)
 
 	e.WriteFileProtocol.Lock()
-	e.WriteFileProtocol.requests[req.MessageData.Id] = &writeMsgResp{
-		ch: respChan,
+	for {
+		if _, ok := e.WriteFileProtocol.requests[req.MessageData.Id]; ok {
+			req.MessageData.Id = uuid.New().String()
+			continue
+		}
+		e.WriteFileProtocol.requests[req.MessageData.Id] = &writeMsgResp{
+			ch: respChan,
+		}
+		break
 	}
 	e.WriteFileProtocol.Unlock()
 
@@ -123,11 +139,27 @@ func (e *protocols) WriteFileAction(id peer.ID, roothash, path string) error {
 			return errors.New(ERR_RespTimeOut)
 		}
 
-		if e.WriteFileProtocol.requests[req.MessageData.Id].WritefileResponse.Code == P2PResponseFinish {
+		e.WriteFileProtocol.Lock()
+		resp, ok := e.WriteFileProtocol.requests[req.MessageData.Id]
+		if !ok {
+			e.WriteFileProtocol.Unlock()
+			return errors.New(ERR_RespFailure)
+		}
+		e.WriteFileProtocol.Unlock()
+
+		if resp.WritefileResponse == nil {
+			return errors.New(ERR_RespFailure)
+		}
+
+		if resp.WritefileResponse.Code == P2PResponseFinish {
 			return nil
 		}
 
-		offset = e.WriteFileProtocol.requests[req.MessageData.Id].WritefileResponse.Offset
+		if resp.WritefileResponse.Code == P2PResponseOK {
+			offset = resp.Offset
+		} else {
+			return errors.New(ERR_RespFailure)
+		}
 	}
 	return errors.New(ERR_RespInvalidData)
 }
@@ -228,7 +260,6 @@ func (e *WriteFileProtocol) onWriteFileRequest(s network.Stream) {
 	} else {
 		resp.Offset = size + int64(data.Length)
 	}
-
 	// send response to the request using the message string he provided
 	e.SendProtoMessage(s.Conn().RemotePeer(), protocol.ID(e.ProtocolPrefix+writeFileResponse), resp)
 }
@@ -262,7 +293,5 @@ func (e *WriteFileProtocol) onWriteFileResponse(s network.Stream) {
 		} else {
 			e.requests[data.MessageData.Id].ch <- false
 		}
-	} else {
-		return
 	}
 }
diff --git a/examples/readfile/example_readfile.go b/examples/readfile/example_readfile.go
@@ -68,7 +68,6 @@ func main() {
 	}
 	h1.Peerstore().AddAddr(info.ID, maddr, time.Hour)
 
-	err = h1.ReadFileAction(info.ID, "roothash", "9e579206ed5b82181d8e5e91a7da261c92ee11094be50a2c2f81a0af846220a9", file, 8388608)
+	err = h1.ReadFileAction(info.ID, "00e5b0a7c4c1d2f89b6347b2c0ea90720e5eb01aa1cf0e9052794383c6afd305", "00e5b0a7c4c1d2f89b6347b2c0ea90720e5eb01aa1cf0e9052794383c6afd305", file, 8388608)
 	fmt.Println("err: ", err)
-	select {}
 }