simplify header handling

This commit is contained in:
xtaci
2018-09-28 19:40:24 +08:00
parent e28719eede
commit b993df6776
4 changed files with 12 additions and 13 deletions
+1 -1
View File
@@ -36,7 +36,7 @@ func newFrame(cmd byte, sid uint32) Frame {
return Frame{ver: version, cmd: cmd, sid: sid} return Frame{ver: version, cmd: cmd, sid: sid}
} }
type rawHeader []byte type rawHeader [headerSize]byte
func (h rawHeader) Version() byte { func (h rawHeader) Version() byte {
return h[0] return h[0]
+10 -10
View File
@@ -210,30 +210,30 @@ func (s *Session) returnTokens(n int) {
// session read a frame from underlying connection // session read a frame from underlying connection
// it's data is pointed to the input buffer // it's data is pointed to the input buffer
func (s *Session) readFrame(buffer []byte) (f Frame, err error) { func (s *Session) readFrame(buffer []byte) (f Frame, err error) {
if _, err := io.ReadFull(s.conn, buffer[:headerSize]); err != nil { var hdr rawHeader
if _, err := io.ReadFull(s.conn, hdr[:]); err != nil {
return f, errors.Wrap(err, "readFrame") return f, errors.Wrap(err, "readFrame")
} }
dec := rawHeader(buffer) if hdr.Version() != version {
if dec.Version() != version {
return f, errors.New(errInvalidProtocol) return f, errors.New(errInvalidProtocol)
} }
f.ver = dec.Version() f.ver = hdr.Version()
f.cmd = dec.Cmd() f.cmd = hdr.Cmd()
f.sid = dec.StreamID() f.sid = hdr.StreamID()
if length := dec.Length(); length > 0 { if length := hdr.Length(); length > 0 {
if _, err := io.ReadFull(s.conn, buffer[headerSize:headerSize+length]); err != nil { f.data = buffer[:length]
if _, err := io.ReadFull(s.conn, f.data); err != nil {
return f, errors.Wrap(err, "readFrame") return f, errors.Wrap(err, "readFrame")
} }
f.data = buffer[headerSize : headerSize+length]
} }
return f, nil return f, nil
} }
// recvLoop keeps on reading from underlying connection if tokens are available // recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() { func (s *Session) recvLoop() {
buffer := make([]byte, (1<<16)+headerSize) buffer := make([]byte, 1<<16)
for { for {
for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() { for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
<-s.bucketNotify <-s.bucketNotify
-1
View File
@@ -523,7 +523,6 @@ func TestRandomFrame(t *testing.T) {
copy(buf[headerSize:], f.data) copy(buf[headerSize:], f.data)
session.conn.Write(buf) session.conn.Write(buf)
t.Log(rawHeader(buf))
cli.Close() cli.Close()
} }
+1 -1
View File
@@ -62,7 +62,7 @@ func (s *Stream) Read(b []byte) (n int, err error) {
READ: READ:
s.bufferLock.Lock() s.bufferLock.Lock()
n, err = s.buffer.Read(b) n, _ = s.buffer.Read(b)
s.bufferLock.Unlock() s.bufferLock.Unlock()
if n > 0 { if n > 0 {