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  }