Golang Conn.WriteProto方法代码示例

func (s *Server) download(conn *transport.Conn, download *pond.Download) *pond.Reply {
	var from [32]byte
	if len(download.From) != len(from) {
		return &pond.Reply{Status: pond.Reply_PARSE_ERROR.Enum()}
	}
	copy(from[:], download.From)

	account, ok := s.getAccount(&from)
	if !ok {
		return &pond.Reply{Status: pond.Reply_NO_SUCH_ADDRESS.Enum()}
	}

	path := filepath.Join(account.FilePath(), strconv.FormatUint(*download.Id, 16))
	file, err := os.OpenFile(path, os.O_RDONLY, 0600)
	if err != nil {
		return &pond.Reply{Status: pond.Reply_NO_SUCH_FILE.Enum()}
	}
	defer file.Close()

	fi, err := file.Stat()
	if err != nil {
		log.Printf("failed to stat file %s: %s", path, err)
		return &pond.Reply{Status: pond.Reply_INTERNAL_ERROR.Enum()}
	}
	size := fi.Size()

	if download.Resume != nil {
		if *download.Resume < 1 {
			return &pond.Reply{Status: pond.Reply_PARSE_ERROR.Enum()}
		}

		if size <= *download.Resume {
			return &pond.Reply{Status: pond.Reply_RESUME_PAST_END_OF_FILE.Enum()}
		}
		pos, err := file.Seek(*download.Resume, 0 /* from start */)
		if pos != *download.Resume || err != nil {
			log.Printf("failed to seek to %d in %s: got %d %s", *download.Resume, path, pos, err)
			return &pond.Reply{Status: pond.Reply_INTERNAL_ERROR.Enum()}
		}
	}

	reply := &pond.Reply{
		Download: &pond.DownloadReply{
			Size: proto.Int64(size),
		},
	}
	if err := conn.WriteProto(reply); err != nil {
		return nil
	}

	io.Copy(conn, file)
	return nil
}

// transferDetachmentConn transfers as much of a detachment as possible on a
// single connection. It calls sendStatus repeatedly with the current state of
// the transfer and watches killChan for an abort signal. It returns an error
// and an indication of whether the error is fatal. If not fatatl then another
// connection can be attempted in order to resume the transfer.
func (c *client) transferDetachmentConn(sendStatus func(s string, done, total int64), conn *transport.Conn, transfer detachmentTransfer, killChan chan bool) (err error, fatal bool) {
	defer conn.Close()

	// transferred is the number of bytes that *this connection* has transferred.
	// total is the full length of the file.
	var startingOffset, transferred, total int64

	sendStatus("Requesting transfer", 0, 0)
	if err := conn.WriteProto(transfer.Request()); err != nil {
		return fmt.Errorf("failed to write request: %s", err), false
	}

	reply := new(pond.Reply)
	if err := conn.ReadProto(reply); err != nil {
		return fmt.Errorf("failed to read reply: %s", err), false
	}

	if reply.Status != nil && *reply.Status == pond.Reply_RESUME_PAST_END_OF_FILE {
		return nil, false
	}

	if err := replyToError(reply); err != nil {
		if reply.GetStatus() == pond.Reply_OVER_QUOTA {
			return fmt.Errorf("server reports that the upload would exceed allowed quota"), true
		}
		return fmt.Errorf("request failed: %s", err), false
	}

	var file *os.File
	var isUpload, isComplete bool
	if file, isUpload, startingOffset, total, isComplete, err = transfer.ProcessReply(reply); err != nil {
		return fmt.Errorf("request failed: %s", err), false
	}
	if isComplete {
		return nil, false
	}
	todo := total - startingOffset

	var in io.Reader
	var out io.Writer
	if isUpload {
		out = conn
		in = file
	} else {
		out = file
		in = conn
	}

	buf := make([]byte, 16*1024)
	var lastUpdate time.Time

	for transferred < todo {
		select {
		case <-killChan:
			return backgroundCanceledError, true
		default:
			break
		}

		conn.SetDeadline(time.Now().Add(30 * time.Second))

		n, err := in.Read(buf)
		if err != nil {
			if isUpload {
				return fmt.Errorf("failed to read from disk: %s", err), true
			}

			return err, false
		}

		n, err = out.Write(buf[:n])
		if err != nil {
			if !isUpload {
				return fmt.Errorf("failed to write to disk: %s", err), true
			}

			return err, false
		}

		transferred += int64(n)

		if transferred > todo {
			return errors.New("transferred more than the expected amount"), true
		}
		now := time.Now()
		if lastUpdate.IsZero() || now.Sub(lastUpdate) > 10*time.Millisecond {
			lastUpdate = now
			sendStatus("", startingOffset+transferred, total)
		}
	}
	sendStatus("", startingOffset+transferred, total)

	if transferred < todo {
		return errors.New("incomplete transfer"), false
	}
//.........这里部分代码省略.........

func (s *Server) upload(from *[32]byte, conn *transport.Conn, upload *pond.Upload) *pond.Reply {
	account, ok := s.getAccount(from)
	if !ok {
		return &pond.Reply{Status: pond.Reply_NO_ACCOUNT.Enum()}
	}

	if *upload.Size < 1 {
		return &pond.Reply{Status: pond.Reply_PARSE_ERROR.Enum()}
	}

	path := filepath.Join(account.FilePath(), strconv.FormatUint(*upload.Id, 16))

	if !account.LoadFileInfo() {
		return &pond.Reply{Status: pond.Reply_INTERNAL_ERROR.Enum()}
	}

	file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0600)
	if err != nil {
		log.Printf("Failed to create file %s: %s", path, err)
		return &pond.Reply{Status: pond.Reply_INTERNAL_ERROR.Enum()}
	}
	defer file.Close()

	offset, err := file.Seek(0, 2 /* from end */)

	switch {
	case offset == *upload.Size:
		return &pond.Reply{Status: pond.Reply_FILE_COMPLETE.Enum()}
	case offset > *upload.Size:
		return &pond.Reply{Status: pond.Reply_FILE_LARGER_THAN_SIZE.Enum()}
	}

	size := *upload.Size - offset
	if !account.ReserveFile(offset > 0, size) {
		return &pond.Reply{Status: pond.Reply_OVER_QUOTA.Enum()}
	}

	var resume *int64
	if offset > 0 {
		resume = proto.Int64(offset)
	}

	reply := &pond.Reply{
		Upload: &pond.UploadReply{
			Resume: resume,
		},
	}
	if err := conn.WriteProto(reply); err != nil {
		return nil
	}

	n, err := io.Copy(file, io.LimitReader(conn, size))
	switch {
	case n == 0:
		os.Remove(path)
		account.ReleaseFile(true, size)
	case n < size:
		account.ReleaseFile(false, size-n)
	case n == size:
		if err == nil {
			conn.Write([]byte{0})
		}
	case n > size:
		panic("impossible")
	}

	return nil
}

func (s *Server) Process(conn *transport.Conn) {
	req := new(pond.Request)
	if err := conn.ReadProto(req); err != nil {
		log.Printf("Error from Read: %s", err)
		return
	}

	from := &conn.Peer
	var reply *pond.Reply
	var messageFetched string

	if req.NewAccount != nil {
		reply = s.newAccount(from, req.NewAccount)
	} else if req.Deliver != nil {
		reply = s.deliver(from, req.Deliver)
	} else if req.Fetch != nil {
		reply, messageFetched = s.fetch(from, req.Fetch)
	} else if req.Upload != nil {
		reply = s.upload(from, conn, req.Upload)
		if reply == nil {
			// Connection will be handled by upload.
			return
		}
	} else if req.Download != nil {
		reply = s.download(conn, req.Download)
		if reply == nil {
			// Connection will be handled by download.
			return
		}
	} else if req.Revocation != nil {
		reply = s.revocation(from, req.Revocation)
	} else {
		reply = &pond.Reply{Status: pond.Reply_NO_REQUEST.Enum()}
	}

	if reply == nil {
		reply = &pond.Reply{}
	}

	if err := conn.WriteProto(reply); err != nil {
		log.Printf("Error from Write: %s", err)
		return
	}

	if err := conn.WaitForClose(); err != nil {
		log.Printf("Error from WaitForClose: %s", err)
		return
	}

	if len(messageFetched) > 0 {
		// We replied to a Fetch and the client successfully acked the
		// message by securely closing the connection. So we can mark
		// the message as delivered.
		s.confirmedDelivery(from, messageFetched)
	}

	s.Lock()
	needSweep := false
	now := time.Now()
	if s.lastSweepTime.IsZero() || now.Before(s.lastSweepTime) || now.Sub(s.lastSweepTime) > sweepInterval {
		s.lastSweepTime = now
		needSweep = true
	}
	s.Unlock()

	if needSweep {
		s.sweep()
	}
}