github.com/oarkflow/log@v1.0.78/gelf/writer.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  	"crypto/rand"
    10  	"encoding/json"
    11  	"fmt"
    12  	"io"
    13  	"net"
    14  )
    15  
    16  type WriterInterface interface {
    17  	Close() error
    18  	Write([]byte) (int, error)
    19  	WriteMessage(*Message) error
    20  }
    21  
    22  // Writer implements io.Writer and is used to send both discrete
    23  // messages to a graylog2 server, or data from a stream-oriented
    24  // interface (like the functions in log).
    25  type Writer struct {
    26  	addr     string
    27  	conn     net.Conn
    28  	hostname string
    29  	Facility string // defaults to current process name
    30  	proto    string
    31  }
    32  
    33  // writes the gzip compressed byte array to the connection as a series
    34  // of GELF chunked messages.  The format is documented at
    35  // http://docs.graylog.org/en/2.1/pages/gelf.html as:
    36  //
    37  //     2-byte magic (0x1e 0x0f), 8 byte id, 1 byte sequence id, 1 byte
    38  //     total, chunk-data
    39  func (w *Writer) writeChunked(zBytes []byte) (err error) {
    40  	b := make([]byte, 0, ChunkSize)
    41  	buf := bytes.NewBuffer(b)
    42  	nChunksI := numChunks(zBytes)
    43  	if nChunksI > 128 {
    44  		return fmt.Errorf("msg too large, would need %d chunks", nChunksI)
    45  	}
    46  	nChunks := uint8(nChunksI)
    47  	// use urandom to get a unique message id
    48  	msgId := make([]byte, 8)
    49  	n, err := io.ReadFull(rand.Reader, msgId)
    50  	if err != nil || n != 8 {
    51  		return fmt.Errorf("rand.Reader: %d/%s", n, err)
    52  	}
    53  
    54  	bytesLeft := len(zBytes)
    55  	for i := uint8(0); i < nChunks; i++ {
    56  		buf.Reset()
    57  		// manually write header.  Don't care about
    58  		// host/network byte order, because the spec only
    59  		// deals in individual bytes.
    60  		buf.Write(magicChunked) //magic
    61  		buf.Write(msgId)
    62  		buf.WriteByte(i)
    63  		buf.WriteByte(nChunks)
    64  		// slice out our chunk from zBytes
    65  		chunkLen := chunkedDataLen
    66  		if chunkLen > bytesLeft {
    67  			chunkLen = bytesLeft
    68  		}
    69  		off := int(i) * chunkedDataLen
    70  		chunk := zBytes[off : off+chunkLen]
    71  		buf.Write(chunk)
    72  
    73  		// write this chunk, and make sure the write was good
    74  		n, err := w.conn.Write(buf.Bytes())
    75  		if err != nil {
    76  			return fmt.Errorf("Write (chunk %d/%d): %s", i,
    77  				nChunks, err)
    78  		}
    79  		if n != len(buf.Bytes()) {
    80  			return fmt.Errorf("Write len: (chunk %d/%d) (%d/%d)",
    81  				i, nChunks, n, len(buf.Bytes()))
    82  		}
    83  
    84  		bytesLeft -= chunkLen
    85  	}
    86  
    87  	if bytesLeft != 0 {
    88  		return fmt.Errorf("error: %d bytes left after sending", bytesLeft)
    89  	}
    90  	return nil
    91  }
    92  
    93  // Close connection and interrupt blocked Read or Write operations
    94  func (w *Writer) Close() error {
    95  	if w.conn == nil {
    96  		return nil
    97  	}
    98  	return w.conn.Close()
    99  }
   100  
   101  func ToByte(in interface{}) []byte {
   102  	o, _ := json.Marshal(in)
   103  	return o
   104  }
   105  
   106  func ToMap(in interface{}) (out map[string]interface{}) {
   107  	o, _ := json.Marshal(in)
   108  	json.Unmarshal(o, &out)
   109  	return
   110  }
   111  
   112  func ByteToMap(in []byte) (out map[string]interface{}) {
   113  	json.Unmarshal(in, &out)
   114  	return
   115  }