out_s3.go

package main

import "github.com/fluent/fluent-bit-go/output"
import "github.com/json-iterator/go"
import "github.com/aws/aws-sdk-go/aws"
import "github.com/aws/aws-sdk-go/aws/awserr"
import "github.com/aws/aws-sdk-go/aws/session"
import "github.com/aws/aws-sdk-go/service/s3"
import "github.com/aws/aws-sdk-go/service/s3/s3manager"
import log "github.com/sirupsen/logrus"
import "github.com/prometheus/common/version"

import (
	"C"
	"bytes"
	"compress/gzip"
	"os"
	"path/filepath"
	"strings"
	"time"
	"unsafe"
)

var plugin GoOutputPlugin = &fluentPlugin{}
var logger *log.Logger

func init() {
	logLevel, _ := log.ParseLevel("info")
	logger = newLogger(logLevel)
}

type s3operator struct {
	bucket         string
	prefix         string
	uploader       *s3manager.Uploader
	compressFormat format
	logger         *log.Logger
	location       *time.Location
}

type GoOutputPlugin interface {
	PluginConfigKey(ctx unsafe.Pointer, key string) string
	Unregister(ctx unsafe.Pointer)
	GetRecord(dec *output.FLBDecoder) (ret int, ts interface{}, rec map[interface{}]interface{})
	NewDecoder(data unsafe.Pointer, length int) *output.FLBDecoder
	Put(s3operator *s3operator, objectKey string, timestamp time.Time, line string) error
	Exit(code int)
}

type fluentPlugin struct{}

func (p *fluentPlugin) PluginConfigKey(ctx unsafe.Pointer, key string) string {
	return output.FLBPluginConfigKey(ctx, key)
}

func (p *fluentPlugin) Unregister(ctx unsafe.Pointer) {
	output.FLBPluginUnregister(ctx)
}

func (p *fluentPlugin) GetRecord(dec *output.FLBDecoder) (int, interface{}, map[interface{}]interface{}) {
	return output.GetRecord(dec)
}

func (p *fluentPlugin) NewDecoder(data unsafe.Pointer, length int) *output.FLBDecoder {
	return output.NewDecoder(data, int(length))
}

func (p *fluentPlugin) Exit(code int) {
	os.Exit(code)
}

func (p *fluentPlugin) Put(s3operator *s3operator, objectKey string, timestamp time.Time, line string) error {
	switch s3operator.compressFormat {
	case plainTextFormat:
		_, err := s3operator.uploader.Upload(&s3manager.UploadInput{
			Bucket: aws.String(s3operator.bucket),
			Key:    aws.String(objectKey),
			Body:   strings.NewReader(line),
		})
		return err
	case gzipFormat:
		compressed, err := makeGzip([]byte(line))
		if err != nil {
			return err
		}
		_, err = s3operator.uploader.Upload(&s3manager.UploadInput{
			Bucket: aws.String(s3operator.bucket),
			Key:    aws.String(objectKey),
			Body:   bytes.NewReader(compressed),
		})
		return err
	}

	return nil
}

// based on https://text.baldanders.info/golang/gzip-operation/
func makeGzip(body []byte) ([]byte, error) {
	var b bytes.Buffer
	err := func() error {
		gw := gzip.NewWriter(&b)
		gw.Name = "fluent-bit-go-s3"
		gw.ModTime = time.Now()

		defer gw.Close()

		if _, err := gw.Write(body); err != nil {
			return err
		}
		return nil
	}()
	return b.Bytes(), err
}

//export FLBPluginRegister
func FLBPluginRegister(ctx unsafe.Pointer) int {
	return output.FLBPluginRegister(ctx, "s3", "S3 Output plugin written in GO!")
}

var (
	s3operators []*s3operator
)

func ensureBucket(session *session.Session, bucket, region *string) (bool, error) {
	svc := s3.New(session)
	var input *s3.CreateBucketInput
	// us-east-1 is default region. So, it needn't specify region in CreateBucketInput.
	if *region == "us-east-1" {
		input = &s3.CreateBucketInput{
			Bucket: bucket,
		}
	} else {
		input = &s3.CreateBucketInput{
			Bucket: bucket,
			CreateBucketConfiguration: &s3.CreateBucketConfiguration{
				LocationConstraint: region,
			},
		}
	}

	result, err := svc.CreateBucket(input)
	logger.Tracef("CreateBucket request result is: %s, err: %s", result, err)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok {
			switch aerr.Code() {
			case s3.ErrCodeBucketAlreadyExists:
				logger.Tracef("Bucket(%s) is already exists.", *bucket)
				return true, nil
			case s3.ErrCodeBucketAlreadyOwnedByYou:
				logger.Tracef("Bucket(%s) is already owned by you.", *bucket)
				return true, nil
			default:
				logger.Tracef("CreateBucket is failed with: %s", aerr.Error())
				return false, aerr
			}
		} else {
			return false, err
		}
	}
	return true, nil
}

func newLogger(logLevel log.Level) *log.Logger {
	logger := log.New()
	logger.Level = logLevel
	return logger
}

func newS3Output(ctx unsafe.Pointer, operatorID int) (*s3operator, error) {
	// Example to retrieve an optional configuration parameter
	credential := plugin.PluginConfigKey(ctx, "Credential")
	accessKeyID := plugin.PluginConfigKey(ctx, "AccessKeyID")
	secretAccessKey := plugin.PluginConfigKey(ctx, "SecretAccessKey")
	bucket := plugin.PluginConfigKey(ctx, "Bucket")
	s3prefix := plugin.PluginConfigKey(ctx, "S3Prefix")
	region := plugin.PluginConfigKey(ctx, "Region")
	compress := plugin.PluginConfigKey(ctx, "Compress")
	endpoint := plugin.PluginConfigKey(ctx, "Endpoint")
	autoCreateBucket := plugin.PluginConfigKey(ctx, "AutoCreateBucket")
	logLevel := plugin.PluginConfigKey(ctx, "LogLevel")
	timeZone := plugin.PluginConfigKey(ctx, "TimeZone")

	config, err := getS3Config(accessKeyID, secretAccessKey, credential, s3prefix, bucket, region, compress, endpoint, autoCreateBucket, logLevel, timeZone)
	if err != nil {
		return nil, err
	}
	logger := newLogger(config.logLevel)

	logger.Infof("[flb-go %d] Starting fluent-bit-go-s3: %v\n", operatorID, version.Info())
	logger.Infof("[flb-go %d] plugin credential parameter = '%s'\n", operatorID, credential)
	logger.Infof("[flb-go %d] plugin accessKeyID parameter = '%s'\n", operatorID, accessKeyID)
	logger.Infof("[flb-go %d] plugin secretAccessKey parameter = '%s'\n", operatorID, secretAccessKey)
	logger.Infof("[flb-go %d] plugin bucket parameter = '%s'\n", operatorID, bucket)
	logger.Infof("[flb-go %d] plugin s3prefix parameter = '%s'\n", operatorID, s3prefix)
	logger.Infof("[flb-go %d] plugin region parameter = '%s'\n", operatorID, region)
	logger.Infof("[flb-go %d] plugin compress parameter = '%s'\n", operatorID, compress)
	logger.Infof("[flb-go %d] plugin endpoint parameter = '%s'\n", operatorID, endpoint)
	logger.Infof("[flb-go %d] plugin autoCreateBucket parameter = '%s'\n", operatorID, autoCreateBucket)
	logger.Infof("[flb-go %d] plugin timeZone parameter = '%s'\n", operatorID, timeZone)

	cfg := aws.Config{
		Credentials: config.credentials,
		Region:      config.region,
	}
	if config.endpoint != "" {
		cfg.WithEndpoint(config.endpoint).WithS3ForcePathStyle(true)
	}

	sess := session.New(&cfg)

	if config.autoCreateBucket == true {
		_, err = ensureBucket(sess, config.bucket, config.region)
		if err != nil {
			return nil, err
		}
	}

	uploader := s3manager.NewUploader(sess, func(u *s3manager.Uploader) {
		u.PartSize = 5 * 1024 * 1024
		u.LeavePartsOnError = true
	})

	s3operator := &s3operator{
		bucket:         *config.bucket,
		prefix:         *config.s3prefix,
		uploader:       uploader,
		compressFormat: config.compress,
		logger:         logger,
		location:       config.location,
	}

	return s3operator, nil

}

func addS3Output(ctx unsafe.Pointer) error {
	operatorID := len(s3operators)
	logger.Infof("[s3operator] id = %d\n", operatorID)
	// Set the context to point to any Go variable
	output.FLBPluginSetContext(ctx, operatorID)
	operator, err := newS3Output(ctx, operatorID)
	if err != nil {
		return err
	}

	s3operators = append(s3operators, operator)
	return nil
}

func getS3Operator(ctx unsafe.Pointer) *s3operator {
	operatorID := output.FLBPluginGetContext(ctx).(int)
	return s3operators[operatorID]
}

//export FLBPluginInit
// (fluentbit will call this)
// ctx (context) pointer to fluentbit context (state/ c code)
func FLBPluginInit(ctx unsafe.Pointer) int {
	err := addS3Output(ctx)
	if err != nil {
		logger.Infof("Error: %s\n", err)
		plugin.Unregister(ctx)
		plugin.Exit(1)
		return output.FLB_ERROR
	}

	return output.FLB_OK
}

//export FLBPluginFlush
func FLBPluginFlush(data unsafe.Pointer, length C.int, tag *C.char) int {
	return output.FLB_OK
}

//export FLBPluginFlushCtx
func FLBPluginFlushCtx(ctx, data unsafe.Pointer, length C.int, tag *C.char) int {
	var ret int
	var record map[interface{}]interface{}

	s3operator := getS3Operator(ctx)
	dec := plugin.NewDecoder(data, int(length))
	var lines string

	for {
		ret, _, record = plugin.GetRecord(dec)
		if ret != 0 {
			break
		}

		line, err := createJSON(record)
		if err != nil {
			s3operator.logger.Warnf("error creating message for S3: %v\n", err)
			continue
		}
		lines += line + "\n"
	}

	objectKey := GenerateObjectKey(s3operator, time.Now())
	err := plugin.Put(s3operator, objectKey, time.Now(), lines)
	if err != nil {
		s3operator.logger.Warnf("error sending message for S3: %v\n", err)
		return output.FLB_RETRY
	}

	// Return options:
	//
	// output.FLB_OK    = data have been processed.
	// output.FLB_ERROR = unrecoverable error, do not try this again.
	// output.FLB_RETRY = retry to flush later.
	return output.FLB_OK
}

// format is S3_PREFIX/S3_TRAILING_PREFIX/date/hour/timestamp_uuid.log
func GenerateObjectKey(s3operator *s3operator, t time.Time) string {
	var fileext string
	switch s3operator.compressFormat {
	case plainTextFormat:
		fileext = ".log"
	case gzipFormat:
		fileext = ".log.gz"
	}
	// Convert time.Time object's Local with specified TimeZone's
	time.Local = s3operator.location
	timestamp := t.Local().Format("20060102150405")
	date := t.Local().Format("20060102")
	hour := t.Local().Format("15")

	fileName := strings.Join([]string{timestamp, fileext}, "")

	objectKey := filepath.Join(s3operator.prefix, date, hour, fileName)
	return objectKey
}

func encodeJSON(record map[interface{}]interface{}) map[string]interface{} {
	m := make(map[string]interface{})

	for k, v := range record {
		switch t := v.(type) {
		case []byte:
			// prevent encoding to base64
			m[k.(string)] = string(t)
		case map[interface{}]interface{}:
			if nextValue, ok := record[k].(map[interface{}]interface{}); ok {
				m[k.(string)] = encodeJSON(nextValue)
			}
		default:
			m[k.(string)] = v
		}
	}

	return m
}

func createJSON(record map[interface{}]interface{}) (string, error) {
	m := encodeJSON(record)

	js, err := jsoniter.Marshal(m)
	if err != nil {
		return "{}", err
	}

	return string(js), nil
}

//export FLBPluginExit
func FLBPluginExit() int {
	return output.FLB_OK
}

func main() {
}