Skip to content

Commit

Permalink
client: seek start 0 before upload
Browse files Browse the repository at this point in the history
  • Loading branch information
@wdvxdr1123
wdvxdr1123 committed Mar 28, 2022
1 parent 4f05838 commit e2a42e5
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 38 deletions.
61 changes: 29 additions & 32 deletions client/internal/highway/bdh.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,13 @@ func (bdh *Transaction) encrypt(key []byte) error {
return nil
}

func (s *Session) UploadBDH(input Transaction) ([]byte, error) {
func (s *Session) UploadBDH(trans Transaction) ([]byte, error) {
if len(s.SsoAddr) == 0 {
return nil, errors.New("srv addrs not found. maybe miss some packet?")
}
addr := s.SsoAddr[0].String()

if err := input.encrypt(s.SessionKey); err != nil {
if err := trans.encrypt(s.SessionKey); err != nil {
return nil, err
}
conn, err := net.DialTimeout("tcp", addr, time.Second*20)
Expand All @@ -59,31 +59,32 @@ func (s *Session) UploadBDH(input Transaction) ([]byte, error) {
const chunkSize = 256 * 1024
var rspExt, chunk []byte
offset := 0
if input.Size > chunkSize {
if trans.Size > chunkSize {
chunk = make([]byte, chunkSize)
} else {
chunk = make([]byte, input.Size)
chunk = make([]byte, trans.Size)
}
for {
rl, err := io.ReadFull(input.Body, chunk)
if errors.Is(err, io.EOF) {
chunk = chunk[:cap(chunk)]
rl, err := io.ReadFull(trans.Body, chunk)
if rl == 0 {
break
}
if errors.Is(err, io.ErrUnexpectedEOF) {
chunk = chunk[:rl]
}
ch := md5.Sum(chunk)
head, _ := proto.Marshal(&pb.ReqDataHighwayHead{
MsgBasehead: s.dataHighwayHead(_REQ_CMD_DATA, 4096, input.CommandID, 2052),
MsgBasehead: s.dataHighwayHead(_REQ_CMD_DATA, 4096, trans.CommandID, 2052),
MsgSeghead: &pb.SegHead{
Filesize: input.Size,
Filesize: trans.Size,
Dataoffset: int64(offset),
Datalength: int32(rl),
Serviceticket: input.Ticket,
Serviceticket: trans.Ticket,
Md5: ch[:],
FileMd5: input.Sum,
FileMd5: trans.Sum,
},
ReqExtendinfo: input.Ext,
ReqExtendinfo: trans.Ext,
})
offset += rl
frame := newFrame(head, chunk)
Expand All @@ -102,25 +103,25 @@ func (s *Session) UploadBDH(input Transaction) ([]byte, error) {
rspExt = rspHead.RspExtendinfo
}
if rspHead.MsgSeghead != nil && rspHead.MsgSeghead.Serviceticket != nil {
input.Ticket = rspHead.MsgSeghead.Serviceticket
trans.Ticket = rspHead.MsgSeghead.Serviceticket
}
}
return rspExt, nil
}

func (s *Session) UploadBDHMultiThread(input Transaction, threadCount int) ([]byte, error) {
func (s *Session) UploadBDHMultiThread(trans Transaction, threadCount int) ([]byte, error) {
// for small file and small thread count,
// use UploadBDH instead of UploadBDHMultiThread
if input.Size < 1024*1024*3 || threadCount < 2 {
return s.UploadBDH(input)
if trans.Size < 1024*1024*3 || threadCount < 2 {
return s.UploadBDH(trans)
}

if len(s.SsoAddr) == 0 {
return nil, errors.New("srv addrs not found. maybe miss some packet?")
}
addr := s.SsoAddr[0].String()

if err := input.encrypt(s.SessionKey); err != nil {
if err := trans.encrypt(s.SessionKey); err != nil {
return nil, err
}

Expand All @@ -130,7 +131,7 @@ func (s *Session) UploadBDHMultiThread(input Transaction, threadCount int) ([]by
completedThread uint32
cond = sync.NewCond(&sync.Mutex{})
offset = int64(0)
count = (input.Size + blockSize - 1) / blockSize
count = (trans.Size + blockSize - 1) / blockSize
id = 0
)
doUpload := func() error {
Expand Down Expand Up @@ -165,31 +166,27 @@ func (s *Session) UploadBDHMultiThread(input Transaction, threadCount int) ([]by
break
}
chunk = chunk[:blockSize]
n, err := io.ReadFull(input.Body, chunk)
n, err := io.ReadFull(trans.Body, chunk)
cond.L.Unlock()

if err != nil {
if err == io.EOF {
break
}
if err == io.ErrUnexpectedEOF {
chunk = chunk[:n]
} else {
return err
}
if n == 0 {
break
}
if errors.Is(err, io.ErrUnexpectedEOF) {
chunk = chunk[:n]
}
ch := md5.Sum(chunk)
head, _ := proto.Marshal(&pb.ReqDataHighwayHead{
MsgBasehead: s.dataHighwayHead(_REQ_CMD_DATA, 4096, input.CommandID, 2052),
MsgBasehead: s.dataHighwayHead(_REQ_CMD_DATA, 4096, trans.CommandID, 2052),
MsgSeghead: &pb.SegHead{
Filesize: input.Size,
Filesize: trans.Size,
Dataoffset: off,
Datalength: int32(n),
Serviceticket: input.Ticket,
Serviceticket: trans.Ticket,
Md5: ch[:],
FileMd5: input.Sum,
FileMd5: trans.Sum,
},
ReqExtendinfo: input.Ext,
ReqExtendinfo: trans.Ext,
})
frame := newFrame(head, chunk)
_, err = frame.WriteTo(conn)
Expand Down
12 changes: 6 additions & 6 deletions client/internal/highway/highway.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ func (s *Session) AppendAddr(ip, port uint32) {
s.SsoAddr = append(s.SsoAddr, addr)
}

func (s *Session) Upload(addr Addr, input Transaction) error {
func (s *Session) Upload(addr Addr, trans Transaction) error {
conn, err := net.DialTimeout("tcp", addr.String(), time.Second*3)
if err != nil {
return errors.Wrap(err, "connect error")
Expand All @@ -60,7 +60,7 @@ func (s *Session) Upload(addr Addr, input Transaction) error {
reader := binary.NewNetworkReader(conn)
for {
chunk = chunk[:chunkSize]
rl, err := io.ReadFull(input.Body, chunk)
rl, err := io.ReadFull(trans.Body, chunk)
if errors.Is(err, io.EOF) {
break
}
Expand All @@ -69,14 +69,14 @@ func (s *Session) Upload(addr Addr, input Transaction) error {
}
ch := md5.Sum(chunk)
head, _ := proto.Marshal(&pb.ReqDataHighwayHead{
MsgBasehead: s.dataHighwayHead(_REQ_CMD_DATA, 4096, input.CommandID, 2052),
MsgBasehead: s.dataHighwayHead(_REQ_CMD_DATA, 4096, trans.CommandID, 2052),
MsgSeghead: &pb.SegHead{
Filesize: input.Size,
Filesize: trans.Size,
Dataoffset: int64(offset),
Datalength: int32(rl),
Serviceticket: input.Ticket,
Serviceticket: trans.Ticket,
Md5: ch[:],
FileMd5: input.Sum,
FileMd5: trans.Sum,
},
ReqExtendinfo: []byte{},
})
Expand Down
2 changes: 2 additions & 0 deletions client/ptt.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,8 @@ func (c *QQClient) UploadShortVideo(target message.Source, video, thumb io.ReadS
cmd = 89
}
ext, _ := proto.Marshal(c.buildPttShortVideoProto(target, videoSum, thumbSum, videoLen, thumbLen).PttShortVideoUploadReq)
_, _ = thumb.Seek(0, io.SeekStart)
_, _ = video.Seek(0, io.SeekStart)
combined := io.MultiReader(thumb, video)
input := highway.Transaction{
CommandID: cmd,
Expand Down

0 comments on commit e2a42e5

Please sign in to comment.