github.com/sujit-baniya/log@v1.0.73/gelf/reader.go (about)

     1  // Copyright 2012 SocialCode. All rights reserved.
     2  // Use of this source code is governed by the MIT
     3  // license that can be found in the LICENSE file.
     4  
     5  package gelf
     6  
     7  import (
     8  	"bytes"
     9  	"compress/gzip"
    10  	"compress/zlib"
    11  	"encoding/json"
    12  	"fmt"
    13  	"io"
    14  	"net"
    15  	"strings"
    16  	"sync"
    17  )
    18  
    19  type Reader struct {
    20  	mu   sync.Mutex
    21  	conn net.Conn
    22  }
    23  
    24  func NewReader(addr string) (*Reader, error) {
    25  	var err error
    26  	udpAddr, err := net.ResolveUDPAddr("udp", addr)
    27  	if err != nil {
    28  		return nil, fmt.Errorf("ResolveUDPAddr('%s'): %s", addr, err)
    29  	}
    30  
    31  	conn, err := net.ListenUDP("udp", udpAddr)
    32  	if err != nil {
    33  		return nil, fmt.Errorf("ListenUDP: %s", err)
    34  	}
    35  
    36  	r := new(Reader)
    37  	r.conn = conn
    38  	return r, nil
    39  }
    40  
    41  func (r *Reader) Addr() string {
    42  	return r.conn.LocalAddr().String()
    43  }
    44  
    45  // FIXME: this will discard data if p isn't big enough to hold the
    46  // full message.
    47  func (r *Reader) Read(p []byte) (int, error) {
    48  	msg, err := r.ReadMessage()
    49  	if err != nil {
    50  		return -1, err
    51  	}
    52  
    53  	var data string
    54  
    55  	if msg.Full == "" {
    56  		data = msg.Short
    57  	} else {
    58  		data = msg.Full
    59  	}
    60  
    61  	return strings.NewReader(data).Read(p)
    62  }
    63  
    64  func (r *Reader) ReadMessage() (*Message, error) {
    65  	cBuf := make([]byte, ChunkSize)
    66  	var (
    67  		err        error
    68  		n, length  int
    69  		cid, ocid  []byte
    70  		seq, total uint8
    71  		cHead      []byte
    72  		cReader    io.Reader
    73  		chunks     [][]byte
    74  	)
    75  
    76  	for got := 0; got < 128 && (total == 0 || got < int(total)); got++ {
    77  		if n, err = r.conn.Read(cBuf); err != nil {
    78  			return nil, fmt.Errorf("Read: %s", err)
    79  		}
    80  		cHead, cBuf = cBuf[:2], cBuf[:n]
    81  
    82  		if bytes.Equal(cHead, magicChunked) {
    83  			//fmt.Printf("chunked %v\n", cBuf[:14])
    84  			cid, seq, total = cBuf[2:2+8], cBuf[2+8], cBuf[2+8+1]
    85  			if ocid != nil && !bytes.Equal(cid, ocid) {
    86  				return nil, fmt.Errorf("out-of-band message %v (awaited %v)", cid, ocid)
    87  			} else if ocid == nil {
    88  				ocid = cid
    89  				chunks = make([][]byte, total)
    90  			}
    91  			n = len(cBuf) - chunkedHeaderLen
    92  			//fmt.Printf("setting chunks[%d]: %d\n", seq, n)
    93  			chunks[seq] = append(make([]byte, 0, n), cBuf[chunkedHeaderLen:]...)
    94  			length += n
    95  		} else { //not chunked
    96  			if total > 0 {
    97  				return nil, fmt.Errorf("out-of-band message (not chunked)")
    98  			}
    99  			break
   100  		}
   101  	}
   102  	//fmt.Printf("\nchunks: %v\n", chunks)
   103  
   104  	if length > 0 {
   105  		if cap(cBuf) < length {
   106  			cBuf = append(cBuf, make([]byte, 0, length-cap(cBuf))...)
   107  		}
   108  		cBuf = cBuf[:0]
   109  		for i := range chunks {
   110  			//fmt.Printf("appending %d %v\n", i, chunks[i])
   111  			cBuf = append(cBuf, chunks[i]...)
   112  		}
   113  		cHead = cBuf[:2]
   114  	}
   115  
   116  	// the data we get from the wire is compressed
   117  	if bytes.Equal(cHead, magicGzip) {
   118  		cReader, err = gzip.NewReader(bytes.NewReader(cBuf))
   119  	} else if cHead[0] == magicZlib[0] &&
   120  		(int(cHead[0])*256+int(cHead[1]))%31 == 0 {
   121  		// zlib is slightly more complicated, but correct
   122  		cReader, err = zlib.NewReader(bytes.NewReader(cBuf))
   123  	} else {
   124  		// compliance with https://github.com/Graylog2/graylog2-server
   125  		// treating all messages as uncompressed if  they are not gzip, zlib or
   126  		// chunked
   127  		cReader = bytes.NewReader(cBuf)
   128  	}
   129  
   130  	if err != nil {
   131  		return nil, fmt.Errorf("NewReader: %s", err)
   132  	}
   133  
   134  	msg := new(Message)
   135  	if err := json.NewDecoder(cReader).Decode(&msg); err != nil {
   136  		return nil, fmt.Errorf("json.Unmarshal: %s", err)
   137  	}
   138  
   139  	return msg, nil
   140  }