github.com/perlchild/DBShield@v0.0.0-20170924200059-c888d9e40e13/dbshield/sql/sql.go (about) 1 package sql 2 3 import ( 4 "bytes" 5 "encoding/binary" 6 "sync" 7 "time" 8 9 "github.com/xwb1989/sqlparser" 10 ) 11 12 //QueryContext holds information around query 13 type QueryContext struct { 14 Query []byte 15 User []byte 16 Client []byte 17 Database []byte 18 Time time.Time 19 } 20 21 //Unmarshal []byte into QueryContext 22 func (c *QueryContext) Unmarshal(b []byte) (size uint32) { 23 n := binary.BigEndian.Uint32(b) 24 b = b[4:] 25 26 c.Query = b[:n] 27 size = n 28 29 b = b[n:] 30 n = binary.BigEndian.Uint32(b) 31 b = b[4:] 32 c.User = b[:n] 33 size += n 34 35 b = b[n:] 36 n = binary.BigEndian.Uint32(b) 37 b = b[4:] 38 c.Client = b[:n] 39 size += n 40 41 b = b[n:] 42 n = binary.BigEndian.Uint32(b) 43 b = b[4:] 44 c.Database = b[:n] 45 size += n 46 47 c.Time.UnmarshalBinary(b[n:]) 48 size += 8 49 return 50 } 51 52 var bufPool = sync.Pool{ 53 New: func() interface{} { 54 return new(bytes.Buffer) 55 }, 56 } 57 58 //Marshal load []byte into QueryContext 59 func (c *QueryContext) Marshal() []byte { 60 buf := bufPool.Get().(*bytes.Buffer) 61 defer bufPool.Put(buf) 62 buf.Reset() 63 l := make([]byte, 4) 64 binary.BigEndian.PutUint32(l, uint32(len(c.Query))) 65 buf.Write(l) 66 buf.Write(c.Query) 67 68 binary.BigEndian.PutUint32(l, uint32(len(c.User))) 69 buf.Write(l) 70 buf.Write(c.User) 71 72 binary.BigEndian.PutUint32(l, uint32(len(c.Client))) 73 buf.Write(l) 74 buf.Write(c.Client) 75 76 binary.BigEndian.PutUint32(l, uint32(len(c.Database))) 77 buf.Write(l) 78 buf.Write(c.Database) 79 80 t, _ := c.Time.MarshalBinary() 81 buf.Write(t) 82 return buf.Bytes() 83 } 84 85 //Pattern returns pattern of given query 86 func Pattern(query []byte) []byte { 87 tokenizer := sqlparser.NewStringTokenizer(string(query)) 88 buf := bytes.Buffer{} 89 l := make([]byte, 4) 90 for { 91 typ, val := tokenizer.Scan() 92 switch typ { 93 case sqlparser.ID: //table, database, variable & ... names 94 buf.Write(val) 95 case 0: //End of query 96 return buf.Bytes() 97 default: 98 binary.BigEndian.PutUint32(l, uint32(typ)) 99 buf.Write(l) 100 } 101 } 102 }