<commit_before>package rabbit

import (
	"net"
	"strings"
	"sync/atomic"
	"time"

	"gopkg.in/tomb.v2"

	"github.com/leandro-lugaresi/message-cannon/runner"
	"github.com/leandro-lugaresi/message-cannon/supervisor"
	"github.com/pkg/errors"
	"github.com/rafaeljesus/retry-go"
	"github.com/speps/go-hashids"
	"github.com/streadway/amqp"
	"go.uber.org/zap"
	"gopkg.in/mcuadros/go-defaults.v1"
)

// Factory is the block responsible for create consumers and restart the rabbitMQ connections.
type Factory struct {
	config Config
	conns  map[string]*amqp.Connection
	log    *zap.Logger
	number int64
}

// NewFactory will open the initial connections and start the recover connections procedure.
func NewFactory(config Config, log *zap.Logger) (*Factory, error) {
	conns := make(map[string]*amqp.Connection)
	for name, cfgConn := range config.Connections {
		defaults.SetDefaults(&cfgConn)
		conn, err := openConnection(cfgConn.DSN, 3, cfgConn.Sleep, cfgConn.Timeout)
		if err != nil {
			return nil, errors.Wrapf(err, "error opening the connection \"%s\"", name)
		}
		conns[name] = conn
	}
	f := &Factory{
		config,
		conns,
		log,
		1,
	}
	return f, nil
}

// CreateConsumers will iterate over config and create all the consumers
func (f *Factory) CreateConsumers() ([]supervisor.Consumer, error) {
	var consumers []supervisor.Consumer
	for name, cfg := range f.config.Consumers {
		defaults.SetDefaults(&cfg)
		consumer, err := f.newConsumer(name, cfg)
		if err != nil {
			return consumers, err
		}
		consumers = append(consumers, consumer)
	}
	return consumers, nil
}

// CreateConsumer create a new consumer for a specific name using the config provided.
func (f *Factory) CreateConsumer(name string) (supervisor.Consumer, error) {
	cfg, ok := f.config.Consumers[name]
	if !ok {
		return nil, errors.Errorf("consumer \"%s\" did not exist", name)
	}
	return f.newConsumer(name, cfg)
}

// Name return the factory name
func (f *Factory) Name() string {
	return "rabbitmq"
}

func (f *Factory) newConsumer(name string, cfg ConsumerConfig) (*consumer, error) {
	_, ok := f.conns[cfg.Connection]
	if !ok {
		available := []string{}
		for cname := range f.conns {
			available = append(available, cname)
		}
		return nil, errors.Errorf(
			"connection name for consumer(%s) did not exist, connections names available: %s",
			name,
			strings.Join(available, ", "))
	}

	f.log.Debug("opening one connection channel", zap.String("connection", cfg.Connection))
	var ch *amqp.Channel
	var errCH error
	conn := f.conns[cfg.Connection]
	ch, errCH = conn.Channel()
	// Reconnect the connection when receive an connection closed error
	if errCH != nil && errCH.Error() == amqp.ErrClosed.Error() {
		cfgConn := f.config.Connections[cfg.Connection]
		conn, err := openConnection(cfgConn.DSN, 5, cfgConn.Sleep, cfgConn.Timeout)
		if err != nil {
			return nil, errors.Wrapf(err, "error reopening the connection \"%s\"", cfg.Connection)
		}
		f.conns[cfg.Connection] = conn
		ch, errCH = conn.Channel()
	}
	if errCH != nil {
		return nil, errors.Wrap(errCH, "failed to open the rabbitMQ channel")
	}
	if len(cfg.DeadLetter) > 0 {
		err := f.declareDeadLetters(ch, cfg.DeadLetter)
		if err != nil {
			return nil, err
		}
	}
	err := f.declareQueue(ch, cfg.Queue)
	if err != nil {
		return nil, err
	}
	f.log.Debug("setting QoS",
		zap.Int("count", cfg.PrefetchCount),
		zap.String("consumer", name))
	if err := ch.Qos(cfg.PrefetchCount, 0, false); err != nil {
		return nil, errors.Wrap(err, "failed to set QoS")
	}
	hash := hashids.New()
	atomic.AddInt64(&f.number, 1)
	hashcounter, err := hash.EncodeInt64([]int64{f.number})
	if err != nil {
		f.log.Warn("Problem generating the hash", zap.Error(err))
	}
	runner, err := runner.New(f.log.With(zap.String("consumer", name)), cfg.Runner)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to create the runner")
	}
	f.log.Info("qtd of workers", zap.Int("workers", cfg.Workers), zap.String("name", name))
	return &consumer{
		queue:       cfg.Queue.Name,
		name:        name,
		hash:        hashcounter,
		opts:        cfg.Options,
		factoryName: f.Name(),
		channel:     ch,
		t:           tomb.Tomb{},
		runner:      runner,
		l:           f.log.With(zap.String("consumer", name)),
		throttle:    make(chan struct{}, cfg.Workers),
		timeout:     cfg.Runner.Timeout,
	}, nil
}

func (f *Factory) declareExchange(ch *amqp.Channel, name string) error {
	if len(name) == 0 {
		f.log.Warn("received a black exchange. wrong config?")
		return nil
	}
	f.log.Debug("declaring an exchange", zap.String("exchange", name))
	ex, ok := f.config.Exchanges[name]
	if !ok {
		f.log.Warn("exchange config didn't exist, we will try to continue")
		return nil
	}
	err := ch.ExchangeDeclare(
		name,
		ex.Type,
		ex.Options.Durable,
		ex.Options.AutoDelete,
		ex.Options.Internal,
		ex.Options.NoWait,
		assertRightTableTypes(ex.Options.Args))
	if err != nil {
		return errors.Wrapf(err, "failed to declare the exchange %s", name)
	}
	return nil
}

func (f *Factory) declareQueue(ch *amqp.Channel, queue QueueConfig) error {
	f.log.Debug("declaring a queue", zap.String("queue", queue.Name))
	q, err := ch.QueueDeclare(
		queue.Name,
		queue.Options.Durable,
		queue.Options.AutoDelete,
		queue.Options.Exclusive,
		queue.Options.NoWait,
		assertRightTableTypes(queue.Options.Args))
	if err != nil {
		return errors.Wrapf(err, "failed to declare the queue \"%s\"", queue.Name)
	}

	for _, b := range queue.Bindings {
		f.log.Debug("adding queue bind",
			zap.String("queue", q.Name),
			zap.String("exchange", b.Exchange))
		err = f.declareExchange(ch, b.Exchange)
		if err != nil {
			return err
		}
		for _, k := range b.RoutingKeys {
			err = ch.QueueBind(q.Name, k, b.Exchange,
				b.Options.NoWait, assertRightTableTypes(b.Options.Args))
			if err != nil {
				return errors.Wrapf(err, "failed to bind the queue \"%s\" to exchange: \"%s\"", q.Name, b.Exchange)
			}
		}
	}
	return nil
}

func (f *Factory) declareDeadLetters(ch *amqp.Channel, name string) error {
	f.log.Debug("declaring a deadletter", zap.String("deadletter", name))
	dead, ok := f.config.DeadLetters[name]
	if !ok {
		f.log.Warn("deadletter config didn't exist, we will try to continue")
		return nil
	}
	err := f.declareQueue(ch, dead.Queue)
	return errors.Wrapf(err, "failed to declare the queue for deadletter %s", name)
}

func openConnection(dsn string, retries int, sleep, timeout time.Duration) (*amqp.Connection, error) {
	var conn *amqp.Connection
	err := retry.Do(func() error {
		var err error
		conn, err = amqp.DialConfig(dsn, amqp.Config{
			Dial: func(network, addr string) (net.Conn, error) {
				return net.DialTimeout(network, addr, timeout)
			},
		})
		return err
	}, retries, sleep)
	return conn, err
}

func assertRightTableTypes(args amqp.Table) amqp.Table {
	for k, v := range args {
		switch v := v.(type) {
		case int:
			args[k] = int64(v)
		}
	}
	return args
}
<commit_msg>fix error with defaults when recreating the consumers<commit_after>package rabbit

import (
	"net"
	"strings"
	"sync/atomic"
	"time"

	"gopkg.in/tomb.v2"

	"github.com/leandro-lugaresi/message-cannon/runner"
	"github.com/leandro-lugaresi/message-cannon/supervisor"
	"github.com/pkg/errors"
	"github.com/rafaeljesus/retry-go"
	"github.com/speps/go-hashids"
	"github.com/streadway/amqp"
	"go.uber.org/zap"
	"gopkg.in/mcuadros/go-defaults.v1"
)

// Factory is the block responsible for create consumers and restart the rabbitMQ connections.
type Factory struct {
	config Config
	conns  map[string]*amqp.Connection
	log    *zap.Logger
	number int64
}

// NewFactory will open the initial connections and start the recover connections procedure.
func NewFactory(config Config, log *zap.Logger) (*Factory, error) {
	conns := make(map[string]*amqp.Connection)
	for name, cfgConn := range config.Connections {
		defaults.SetDefaults(&cfgConn)
		conn, err := openConnection(cfgConn.DSN, 3, cfgConn.Sleep, cfgConn.Timeout)
		if err != nil {
			return nil, errors.Wrapf(err, "error opening the connection \"%s\"", name)
		}
		conns[name] = conn
	}
	f := &Factory{
		config,
		conns,
		log,
		1,
	}
	return f, nil
}

// CreateConsumers will iterate over config and create all the consumers
func (f *Factory) CreateConsumers() ([]supervisor.Consumer, error) {
	var consumers []supervisor.Consumer
	for name, cfg := range f.config.Consumers {
		consumer, err := f.newConsumer(name, cfg)
		if err != nil {
			return consumers, err
		}
		consumers = append(consumers, consumer)
	}
	return consumers, nil
}

// CreateConsumer create a new consumer for a specific name using the config provided.
func (f *Factory) CreateConsumer(name string) (supervisor.Consumer, error) {
	cfg, ok := f.config.Consumers[name]
	if !ok {
		return nil, errors.Errorf("consumer \"%s\" did not exist", name)
	}
	return f.newConsumer(name, cfg)
}

// Name return the factory name
func (f *Factory) Name() string {
	return "rabbitmq"
}

func (f *Factory) newConsumer(name string, cfg ConsumerConfig) (*consumer, error) {
	defaults.SetDefaults(&cfg)
	_, ok := f.conns[cfg.Connection]
	if !ok {
		available := []string{}
		for cname := range f.conns {
			available = append(available, cname)
		}
		return nil, errors.Errorf(
			"connection name for consumer(%s) did not exist, connections names available: %s",
			name,
			strings.Join(available, ", "))
	}

	f.log.Debug("opening one connection channel", zap.String("connection", cfg.Connection))
	var ch *amqp.Channel
	var errCH error
	conn := f.conns[cfg.Connection]
	ch, errCH = conn.Channel()
	// Reconnect the connection when receive an connection closed error
	if errCH != nil && errCH.Error() == amqp.ErrClosed.Error() {
		cfgConn := f.config.Connections[cfg.Connection]
		conn, err := openConnection(cfgConn.DSN, 5, cfgConn.Sleep, cfgConn.Timeout)
		if err != nil {
			return nil, errors.Wrapf(err, "error reopening the connection \"%s\"", cfg.Connection)
		}
		f.conns[cfg.Connection] = conn
		ch, errCH = conn.Channel()
	}
	if errCH != nil {
		return nil, errors.Wrap(errCH, "failed to open the rabbitMQ channel")
	}
	if len(cfg.DeadLetter) > 0 {
		err := f.declareDeadLetters(ch, cfg.DeadLetter)
		if err != nil {
			return nil, err
		}
	}
	err := f.declareQueue(ch, cfg.Queue)
	if err != nil {
		return nil, err
	}
	f.log.Debug("setting QoS",
		zap.Int("count", cfg.PrefetchCount),
		zap.String("consumer", name))
	if err := ch.Qos(cfg.PrefetchCount, 0, false); err != nil {
		return nil, errors.Wrap(err, "failed to set QoS")
	}
	hash := hashids.New()
	atomic.AddInt64(&f.number, 1)
	hashcounter, err := hash.EncodeInt64([]int64{f.number})
	if err != nil {
		f.log.Warn("Problem generating the hash", zap.Error(err))
	}
	runner, err := runner.New(f.log.With(zap.String("consumer", name)), cfg.Runner)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to create the runner")
	}
	f.log.Info("qtd of workers", zap.Int("workers", cfg.Workers), zap.String("name", name))
	return &consumer{
		queue:       cfg.Queue.Name,
		name:        name,
		hash:        hashcounter,
		opts:        cfg.Options,
		factoryName: f.Name(),
		channel:     ch,
		t:           tomb.Tomb{},
		runner:      runner,
		l:           f.log.With(zap.String("consumer", name)),
		throttle:    make(chan struct{}, cfg.Workers),
		timeout:     cfg.Runner.Timeout,
	}, nil
}

func (f *Factory) declareExchange(ch *amqp.Channel, name string) error {
	if len(name) == 0 {
		f.log.Warn("received a black exchange. wrong config?")
		return nil
	}
	f.log.Debug("declaring an exchange", zap.String("exchange", name))
	ex, ok := f.config.Exchanges[name]
	if !ok {
		f.log.Warn("exchange config didn't exist, we will try to continue")
		return nil
	}
	err := ch.ExchangeDeclare(
		name,
		ex.Type,
		ex.Options.Durable,
		ex.Options.AutoDelete,
		ex.Options.Internal,
		ex.Options.NoWait,
		assertRightTableTypes(ex.Options.Args))
	if err != nil {
		return errors.Wrapf(err, "failed to declare the exchange %s", name)
	}
	return nil
}

func (f *Factory) declareQueue(ch *amqp.Channel, queue QueueConfig) error {
	f.log.Debug("declaring a queue", zap.String("queue", queue.Name))
	q, err := ch.QueueDeclare(
		queue.Name,
		queue.Options.Durable,
		queue.Options.AutoDelete,
		queue.Options.Exclusive,
		queue.Options.NoWait,
		assertRightTableTypes(queue.Options.Args))
	if err != nil {
		return errors.Wrapf(err, "failed to declare the queue \"%s\"", queue.Name)
	}

	for _, b := range queue.Bindings {
		f.log.Debug("adding queue bind",
			zap.String("queue", q.Name),
			zap.String("exchange", b.Exchange))
		err = f.declareExchange(ch, b.Exchange)
		if err != nil {
			return err
		}
		for _, k := range b.RoutingKeys {
			err = ch.QueueBind(q.Name, k, b.Exchange,
				b.Options.NoWait, assertRightTableTypes(b.Options.Args))
			if err != nil {
				return errors.Wrapf(err, "failed to bind the queue \"%s\" to exchange: \"%s\"", q.Name, b.Exchange)
			}
		}
	}
	return nil
}

func (f *Factory) declareDeadLetters(ch *amqp.Channel, name string) error {
	f.log.Debug("declaring a deadletter", zap.String("deadletter", name))
	dead, ok := f.config.DeadLetters[name]
	if !ok {
		f.log.Warn("deadletter config didn't exist, we will try to continue")
		return nil
	}
	err := f.declareQueue(ch, dead.Queue)
	return errors.Wrapf(err, "failed to declare the queue for deadletter %s", name)
}

func openConnection(dsn string, retries int, sleep, timeout time.Duration) (*amqp.Connection, error) {
	var conn *amqp.Connection
	err := retry.Do(func() error {
		var err error
		conn, err = amqp.DialConfig(dsn, amqp.Config{
			Dial: func(network, addr string) (net.Conn, error) {
				return net.DialTimeout(network, addr, timeout)
			},
		})
		return err
	}, retries, sleep)
	return conn, err
}

func assertRightTableTypes(args amqp.Table) amqp.Table {
	for k, v := range args {
		switch v := v.(type) {
		case int:
			args[k] = int64(v)
		}
	}
	return args
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package grafeas

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"testing"

	"google.golang.org/grpc/credentials"

	"github.com/golang/protobuf/ptypes/empty"
	kritisv1beta1 "github.com/grafeas/kritis/pkg/kritis/apis/kritis/v1beta1"
	"github.com/grafeas/kritis/pkg/kritis/secrets"
	"github.com/grafeas/kritis/pkg/kritis/testutil"
	"google.golang.org/genproto/googleapis/devtools/containeranalysis/v1beta1/grafeas"
	"google.golang.org/grpc"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var (
	// Certificates valid for 10 years
	// Generated using cfssl
	clientCert = filepath.Join("testdata", "client.pem")
	clientKey  = filepath.Join("testdata", "client-key.pem")
	serverCert = filepath.Join("testdata", "server.pem")
	serverKey  = filepath.Join("testdata", "server-key.pem")
	ca         = filepath.Join("testdata", "ca.pem")
)

func TestNewClientTLS(t *testing.T) {
	keyPair, err := tls.LoadX509KeyPair(serverCert, serverKey)
	if err != nil {
		t.Fatal(err)
	}
	if err != nil {
		t.Fatal(err)
	}

	caData, err := ioutil.ReadFile(ca)
	if err != nil {
		t.Fatal(err)
	}
	caCertPool := x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caData)
	config := kritisv1beta1.GrafeasConfigSpec{Addr: "127.0.0.1:9995", ClientKeyPath: clientKey, ClientCertPath: clientCert, CAPath: ca}
	lis, err := net.Listen("tcp", config.Addr)
	if err != nil {
		t.Fatal(err)
	}
	creds := credentials.NewTLS(&tls.Config{
		Certificates: []tls.Certificate{keyPair},
		ClientCAs:    caCertPool,
		ClientAuth:   tls.RequireAndVerifyClientCert,
	})

	gs := grpc.NewServer(grpc.Creds(creds))
	grafeasMock := newGrafeasServerMock()
	grafeas.RegisterGrafeasV1Beta1Server(gs, grafeasMock)
	go func() {
		if err := gs.Serve(lis); err != nil {
			t.Fatal(err)
		}
	}()
	defer func() {
		gs.GracefulStop()
	}()
	client, err := New(config)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := client.client.CreateOccurrence(context.Background(), &grafeas.CreateOccurrenceRequest{
		Occurrence: &grafeas.Occurrence{},
	}); err != nil {
		t.Fatal(err)
	}
}

func TestValidateConfig(t *testing.T) {
	for _, tt := range []struct {
		config      kritisv1beta1.GrafeasConfigSpec
		expectedErr bool
	}{
		{config: kritisv1beta1.GrafeasConfigSpec{Addr: "/socketaddr"}, expectedErr: false},
		// Missing certificates
		{config: kritisv1beta1.GrafeasConfigSpec{Addr: "addr:port"}, expectedErr: true},
		// Missing address
		{config: kritisv1beta1.GrafeasConfigSpec{}, expectedErr: true},
		// Missing files
		{config: kritisv1beta1.GrafeasConfigSpec{Addr: "host:port", CAPath: "ca", ClientKeyPath: "clientcert", ClientCertPath: "path"}, expectedErr: true},
		{config: kritisv1beta1.GrafeasConfigSpec{Addr: "host:port", CAPath: ca, ClientKeyPath: clientKey, ClientCertPath: clientCert}, expectedErr: false},
	} {
		err := ValidateConfig(tt.config)
		if err != nil && !tt.expectedErr {
			t.Fatalf("Expected no error but got %v", err)
		} else if err == nil && tt.expectedErr {
			t.Fatalf("Expected error but got none")
		}
	}
}

func TestCreateAttestationNoteAndOccurrence(t *testing.T) {
	socketPath, err := filepath.Abs(".grafeas.sock")
	if err != nil {
		t.Fatal(err)
	}
	server := grpc.NewServer()
	grafeasMock := newGrafeasServerMock()
	lis, err := net.Listen("unix", socketPath)
	if err != nil {
		t.Fatalf("Failed to listen on socket %v", err)
	}
	_, err = os.Stat(socketPath)
	t.Logf("File %v", err)
	grafeas.RegisterGrafeasV1Beta1Server(server, grafeasMock)
	go func() {
		if err := server.Serve(lis); err != nil {
			t.Fatalf("Failed to server %v", err)
		}
	}()
	defer func() {
		server.GracefulStop()
		os.Remove(socketPath)
	}()
	client, err := New(kritisv1beta1.GrafeasConfigSpec{Addr: socketPath})
	if err != nil {
		t.Fatalf("Could not initialize the client %v", err)
	}
	aa := &kritisv1beta1.AttestationAuthority{
		Spec: kritisv1beta1.AttestationAuthoritySpec{
			NoteReference: fmt.Sprintf("%s/projects/%s", "api", DefaultProject),
		},
		ObjectMeta: metav1.ObjectMeta{
			Name: "note1",
		},
	}
	if _, err := client.CreateAttestationNote(aa); err != nil {
		t.Fatalf("Unexpected error while creating Note %v", err)
	}
	note, err := client.AttestationNote(aa)
	if err != nil {
		t.Fatalf("Unexpected no error while getting attestation note %v", err)
	}
	expectedNoteName := fmt.Sprintf("projects/%s/notes/note1", DefaultProject)
	if note.Name != expectedNoteName {
		t.Fatalf("Expected %s.\n Got %s", expectedNoteName, note.Name)
	}
	actualHint := note.GetAttestationAuthority().Hint.GetHumanReadableName()
	if actualHint != "note1" {
		t.Fatalf("Expected %s.\n Got %s", expectedNoteName, actualHint)
	}
	// Test Create Attestation Occurrence
	pub, priv := testutil.CreateKeyPair(t, "test")
	pgpKey, err := secrets.NewPgpKey(priv, "", pub)
	if err != nil {
		t.Fatalf("Unexpected error while creating PGP key %v", err)
	}
	secret := &secrets.PGPSigningSecret{
		PgpKey:     pgpKey,
		SecretName: "test",
	}
	occ, err := client.CreateAttestationOccurence(note, testutil.IntTestImage, secret)
	if err != nil {
		t.Fatalf("Unexpected error while creating Occurence %v", err)
	}
	expectedPgpKeyID := pgpKey.Fingerprint()
	if err != nil {
		t.Fatalf("Unexpected error while extracting PGP key id %v", err)
	}

	pgpKeyID := occ.GetAttestation().GetAttestation().GetPgpSignedAttestation().GetPgpKeyId()
	if pgpKeyID != expectedPgpKeyID {
		t.Errorf("Expected PGP key id: %q, got %q", expectedPgpKeyID, pgpKeyID)
	}
	occurrences, err := client.Attestations(testutil.IntTestImage)
	if err != nil {
		t.Fatalf("Unexpected error while listing Occ %v", err)
	}
	if occurrences == nil {
		t.Fatal("Should have created at least 1 occurrence")
	}
}

// grafeasServerMock is a mock for grafeas grpc API
type grafeasServerMock struct {
	notes       map[string]*grafeas.Note
	occurrences map[string]*grafeas.Occurrence
}

func newGrafeasServerMock() *grafeasServerMock {
	return &grafeasServerMock{notes: make(map[string]*grafeas.Note), occurrences: make(map[string]*grafeas.Occurrence)}
}

func (g *grafeasServerMock) GetOccurrence(ctx context.Context, req *grafeas.GetOccurrenceRequest) (*grafeas.Occurrence, error) {
	return g.occurrences[req.Name], nil
}

func (g *grafeasServerMock) ListOccurrences(context.Context, *grafeas.ListOccurrencesRequest) (*grafeas.ListOccurrencesResponse, error) {
	var resp grafeas.ListOccurrencesResponse
	for _, occ := range g.occurrences {
		resp.Occurrences = append(resp.Occurrences, occ)
	}
	return &resp, nil
}

func (g *grafeasServerMock) DeleteOccurrence(context.Context, *grafeas.DeleteOccurrenceRequest) (*empty.Empty, error) {
	return nil, nil
}

func (g *grafeasServerMock) CreateOccurrence(ctx context.Context, req *grafeas.CreateOccurrenceRequest) (*grafeas.Occurrence, error) {
	occ := req.Occurrence
	g.occurrences[occ.Name] = occ
	return occ, nil
}

func (g *grafeasServerMock) BatchCreateOccurrences(context.Context, *grafeas.BatchCreateOccurrencesRequest) (*grafeas.BatchCreateOccurrencesResponse, error) {
	return nil, nil
}

func (g *grafeasServerMock) UpdateOccurrence(context.Context, *grafeas.UpdateOccurrenceRequest) (*grafeas.Occurrence, error) {
	return nil, nil
}

func (g *grafeasServerMock) GetOccurrenceNote(context.Context, *grafeas.GetOccurrenceNoteRequest) (*grafeas.Note, error) {
	return nil, nil
}

func (g *grafeasServerMock) GetNote(ctx context.Context, req *grafeas.GetNoteRequest) (*grafeas.Note, error) {
	return g.notes[req.Name], nil
}

func (g *grafeasServerMock) ListNotes(context.Context, *grafeas.ListNotesRequest) (*grafeas.ListNotesResponse, error) {
	return nil, nil
}

func (g *grafeasServerMock) DeleteNote(context.Context, *grafeas.DeleteNoteRequest) (*empty.Empty, error) {
	return nil, nil
}

func (g *grafeasServerMock) CreateNote(ctx context.Context, req *grafeas.CreateNoteRequest) (*grafeas.Note, error) {
	note := req.Note
	g.notes[note.Name] = note
	return note, nil
}

func (g *grafeasServerMock) BatchCreateNotes(context.Context, *grafeas.BatchCreateNotesRequest) (*grafeas.BatchCreateNotesResponse, error) {
	return nil, nil
}

func (g *grafeasServerMock) UpdateNote(context.Context, *grafeas.UpdateNoteRequest) (*grafeas.Note, error) {
	return nil, nil
}

func (g *grafeasServerMock) ListNoteOccurrences(context.Context, *grafeas.ListNoteOccurrencesRequest) (*grafeas.ListNoteOccurrencesResponse, error) {
	return nil, nil
}

func (g *grafeasServerMock) GetVulnerabilityOccurrencesSummary(context.Context, *grafeas.GetVulnerabilityOccurrencesSummaryRequest) (*grafeas.VulnerabilityOccurrencesSummary, error) {
	return nil, nil
}
<commit_msg>removed unhelpful log line<commit_after>/*
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package grafeas

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"testing"

	"google.golang.org/grpc/credentials"

	"github.com/golang/protobuf/ptypes/empty"
	kritisv1beta1 "github.com/grafeas/kritis/pkg/kritis/apis/kritis/v1beta1"
	"github.com/grafeas/kritis/pkg/kritis/secrets"
	"github.com/grafeas/kritis/pkg/kritis/testutil"
	"google.golang.org/genproto/googleapis/devtools/containeranalysis/v1beta1/grafeas"
	"google.golang.org/grpc"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var (
	// Certificates valid for 10 years
	// Generated using cfssl
	clientCert = filepath.Join("testdata", "client.pem")
	clientKey  = filepath.Join("testdata", "client-key.pem")
	serverCert = filepath.Join("testdata", "server.pem")
	serverKey  = filepath.Join("testdata", "server-key.pem")
	ca         = filepath.Join("testdata", "ca.pem")
)

func TestNewClientTLS(t *testing.T) {
	keyPair, err := tls.LoadX509KeyPair(serverCert, serverKey)
	if err != nil {
		t.Fatal(err)
	}
	if err != nil {
		t.Fatal(err)
	}

	caData, err := ioutil.ReadFile(ca)
	if err != nil {
		t.Fatal(err)
	}
	caCertPool := x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caData)
	config := kritisv1beta1.GrafeasConfigSpec{Addr: "127.0.0.1:9995", ClientKeyPath: clientKey, ClientCertPath: clientCert, CAPath: ca}
	lis, err := net.Listen("tcp", config.Addr)
	if err != nil {
		t.Fatal(err)
	}
	creds := credentials.NewTLS(&tls.Config{
		Certificates: []tls.Certificate{keyPair},
		ClientCAs:    caCertPool,
		ClientAuth:   tls.RequireAndVerifyClientCert,
	})

	gs := grpc.NewServer(grpc.Creds(creds))
	grafeasMock := newGrafeasServerMock()
	grafeas.RegisterGrafeasV1Beta1Server(gs, grafeasMock)
	go func() {
		if err := gs.Serve(lis); err != nil {
			t.Fatal(err)
		}
	}()
	defer func() {
		gs.GracefulStop()
	}()
	client, err := New(config)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := client.client.CreateOccurrence(context.Background(), &grafeas.CreateOccurrenceRequest{
		Occurrence: &grafeas.Occurrence{},
	}); err != nil {
		t.Fatal(err)
	}
}

func TestValidateConfig(t *testing.T) {
	for _, tt := range []struct {
		config      kritisv1beta1.GrafeasConfigSpec
		expectedErr bool
	}{
		{config: kritisv1beta1.GrafeasConfigSpec{Addr: "/socketaddr"}, expectedErr: false},
		// Missing certificates
		{config: kritisv1beta1.GrafeasConfigSpec{Addr: "addr:port"}, expectedErr: true},
		// Missing address
		{config: kritisv1beta1.GrafeasConfigSpec{}, expectedErr: true},
		// Missing files
		{config: kritisv1beta1.GrafeasConfigSpec{Addr: "host:port", CAPath: "ca", ClientKeyPath: "clientcert", ClientCertPath: "path"}, expectedErr: true},
		{config: kritisv1beta1.GrafeasConfigSpec{Addr: "host:port", CAPath: ca, ClientKeyPath: clientKey, ClientCertPath: clientCert}, expectedErr: false},
	} {
		err := ValidateConfig(tt.config)
		if err != nil && !tt.expectedErr {
			t.Fatalf("Expected no error but got %v", err)
		} else if err == nil && tt.expectedErr {
			t.Fatalf("Expected error but got none")
		}
	}
}

func TestCreateAttestationNoteAndOccurrence(t *testing.T) {
	socketPath, err := filepath.Abs(".grafeas.sock")
	if err != nil {
		t.Fatal(err)
	}
	server := grpc.NewServer()
	grafeasMock := newGrafeasServerMock()
	lis, err := net.Listen("unix", socketPath)
	if err != nil {
		t.Fatalf("Failed to listen on socket %v", err)
	}
	_, err = os.Stat(socketPath)
	grafeas.RegisterGrafeasV1Beta1Server(server, grafeasMock)
	go func() {
		if err := server.Serve(lis); err != nil {
			t.Fatalf("Failed to server %v", err)
		}
	}()
	defer func() {
		server.GracefulStop()
		os.Remove(socketPath)
	}()
	client, err := New(kritisv1beta1.GrafeasConfigSpec{Addr: socketPath})
	if err != nil {
		t.Fatalf("Could not initialize the client %v", err)
	}
	aa := &kritisv1beta1.AttestationAuthority{
		Spec: kritisv1beta1.AttestationAuthoritySpec{
			NoteReference: fmt.Sprintf("%s/projects/%s", "api", DefaultProject),
		},
		ObjectMeta: metav1.ObjectMeta{
			Name: "note1",
		},
	}
	if _, err := client.CreateAttestationNote(aa); err != nil {
		t.Fatalf("Unexpected error while creating Note %v", err)
	}
	note, err := client.AttestationNote(aa)
	if err != nil {
		t.Fatalf("Unexpected no error while getting attestation note %v", err)
	}
	expectedNoteName := fmt.Sprintf("projects/%s/notes/note1", DefaultProject)
	if note.Name != expectedNoteName {
		t.Fatalf("Expected %s.\n Got %s", expectedNoteName, note.Name)
	}
	actualHint := note.GetAttestationAuthority().Hint.GetHumanReadableName()
	if actualHint != "note1" {
		t.Fatalf("Expected %s.\n Got %s", expectedNoteName, actualHint)
	}
	// Test Create Attestation Occurrence
	pub, priv := testutil.CreateKeyPair(t, "test")
	pgpKey, err := secrets.NewPgpKey(priv, "", pub)
	if err != nil {
		t.Fatalf("Unexpected error while creating PGP key %v", err)
	}
	secret := &secrets.PGPSigningSecret{
		PgpKey:     pgpKey,
		SecretName: "test",
	}
	occ, err := client.CreateAttestationOccurence(note, testutil.IntTestImage, secret)
	if err != nil {
		t.Fatalf("Unexpected error while creating Occurence %v", err)
	}
	expectedPgpKeyID := pgpKey.Fingerprint()
	if err != nil {
		t.Fatalf("Unexpected error while extracting PGP key id %v", err)
	}

	pgpKeyID := occ.GetAttestation().GetAttestation().GetPgpSignedAttestation().GetPgpKeyId()
	if pgpKeyID != expectedPgpKeyID {
		t.Errorf("Expected PGP key id: %q, got %q", expectedPgpKeyID, pgpKeyID)
	}
	occurrences, err := client.Attestations(testutil.IntTestImage)
	if err != nil {
		t.Fatalf("Unexpected error while listing Occ %v", err)
	}
	if occurrences == nil {
		t.Fatal("Should have created at least 1 occurrence")
	}
}

// grafeasServerMock is a mock for grafeas grpc API
type grafeasServerMock struct {
	notes       map[string]*grafeas.Note
	occurrences map[string]*grafeas.Occurrence
}

func newGrafeasServerMock() *grafeasServerMock {
	return &grafeasServerMock{notes: make(map[string]*grafeas.Note), occurrences: make(map[string]*grafeas.Occurrence)}
}

func (g *grafeasServerMock) GetOccurrence(ctx context.Context, req *grafeas.GetOccurrenceRequest) (*grafeas.Occurrence, error) {
	return g.occurrences[req.Name], nil
}

func (g *grafeasServerMock) ListOccurrences(context.Context, *grafeas.ListOccurrencesRequest) (*grafeas.ListOccurrencesResponse, error) {
	var resp grafeas.ListOccurrencesResponse
	for _, occ := range g.occurrences {
		resp.Occurrences = append(resp.Occurrences, occ)
	}
	return &resp, nil
}

func (g *grafeasServerMock) DeleteOccurrence(context.Context, *grafeas.DeleteOccurrenceRequest) (*empty.Empty, error) {
	return nil, nil
}

func (g *grafeasServerMock) CreateOccurrence(ctx context.Context, req *grafeas.CreateOccurrenceRequest) (*grafeas.Occurrence, error) {
	occ := req.Occurrence
	g.occurrences[occ.Name] = occ
	return occ, nil
}

func (g *grafeasServerMock) BatchCreateOccurrences(context.Context, *grafeas.BatchCreateOccurrencesRequest) (*grafeas.BatchCreateOccurrencesResponse, error) {
	return nil, nil
}

func (g *grafeasServerMock) UpdateOccurrence(context.Context, *grafeas.UpdateOccurrenceRequest) (*grafeas.Occurrence, error) {
	return nil, nil
}

func (g *grafeasServerMock) GetOccurrenceNote(context.Context, *grafeas.GetOccurrenceNoteRequest) (*grafeas.Note, error) {
	return nil, nil
}

func (g *grafeasServerMock) GetNote(ctx context.Context, req *grafeas.GetNoteRequest) (*grafeas.Note, error) {
	return g.notes[req.Name], nil
}

func (g *grafeasServerMock) ListNotes(context.Context, *grafeas.ListNotesRequest) (*grafeas.ListNotesResponse, error) {
	return nil, nil
}

func (g *grafeasServerMock) DeleteNote(context.Context, *grafeas.DeleteNoteRequest) (*empty.Empty, error) {
	return nil, nil
}

func (g *grafeasServerMock) CreateNote(ctx context.Context, req *grafeas.CreateNoteRequest) (*grafeas.Note, error) {
	note := req.Note
	g.notes[note.Name] = note
	return note, nil
}

func (g *grafeasServerMock) BatchCreateNotes(context.Context, *grafeas.BatchCreateNotesRequest) (*grafeas.BatchCreateNotesResponse, error) {
	return nil, nil
}

func (g *grafeasServerMock) UpdateNote(context.Context, *grafeas.UpdateNoteRequest) (*grafeas.Note, error) {
	return nil, nil
}

func (g *grafeasServerMock) ListNoteOccurrences(context.Context, *grafeas.ListNoteOccurrencesRequest) (*grafeas.ListNoteOccurrencesResponse, error) {
	return nil, nil
}

func (g *grafeasServerMock) GetVulnerabilityOccurrencesSummary(context.Context, *grafeas.GetVulnerabilityOccurrencesSummaryRequest) (*grafeas.VulnerabilityOccurrencesSummary, error) {
	return nil, nil
}
<|endoftext|>


<commit_before>package java

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/ezbuy/tgen/global"
	"github.com/ezbuy/tgen/langs"
	"github.com/ezbuy/tgen/tmpl"
	"github.com/samuel/go-thrift/parser"
)

const (
	javaLang = "java"
	javaExt  = ".java"
)

const (
	JavaTypeshort  = "short"
	JavaTypeint    = "int"
	JavaTypelong   = "long"
	JavaTypebool   = "boolean"
	JavaTypebyte   = "byte"
	JavaTypedouble = "double"

	JavaTypeString = "String"

	JavaTypeShort  = "Short"
	JavaTypeInt    = "Integer"
	JavaTypeLong   = "Long"
	JavaTypeBool   = "Boolean"
	JavaTypeByte   = "Byte"
	JavaTypeDouble = "Double"

	// other types (such as array, map, etc.) are implemented in the method 'Typecast'
)

const (
	TPL_ENUM    = "tgen/java/enum"
	TPL_STRUCT  = "tgen/java/struct"
	TPL_SERVICE = "tgen/java/service"
)

var plaintypemapping = map[string]string{
	langs.ThriftTypeI16:    JavaTypeshort,
	langs.ThriftTypeI32:    JavaTypeint,
	langs.ThriftTypeI64:    JavaTypelong,
	langs.ThriftTypeString: JavaTypeString,
	langs.ThriftTypeByte:   JavaTypebyte,
	langs.ThriftTypeBool:   JavaTypebool,
	langs.ThriftTypeDouble: JavaTypedouble,
}

var objecttypemapping = map[string]string{
	langs.ThriftTypeI16:    JavaTypeShort,
	langs.ThriftTypeI32:    JavaTypeInt,
	langs.ThriftTypeI64:    JavaTypeLong,
	langs.ThriftTypeString: JavaTypeString,
	langs.ThriftTypeByte:   JavaTypeByte,
	langs.ThriftTypeBool:   JavaTypeBool,
	langs.ThriftTypeDouble: JavaTypeDouble,
}

type JavaGen struct {
	langs.BaseGen
}

type BaseJava struct {
	Namespace string
	t         *parser.Thrift
	ts        map[string]*parser.Thrift
}

func (b *BaseJava) FilterVariableName(n string) string {
	if b.IsKeyword(n) {
		return fmt.Sprintf("t%s%s", strings.ToUpper(n[:1]), n[1:])
	}
	return n
}

func (b *BaseJava) IsKeyword(n string) bool {
	switch n {
	case "package", "int", "short", "long", "byte", "boolean", "case", "switch", "if ", "for", "else",
		"goto", "Integer", "Short", "Long", "Byte", "Boolean", "class", "break", "try", "catch",
		"double", "Double", "do", "while", "final", "finally", "continue", "interface", "private",
		"public", "protected", "return", "this", "throw", "static", "super", "throws",
		"true", "false", "float", "volatile", "synchronized", "abstract", "default", "extends",
		"native", "new":
		return true
	}
	return false
}

func (b *BaseJava) PlainTypecast(t *parser.Type) string {
	return b.typecast(t, true)
}

func (b *BaseJava) ObjectTypecast(t *parser.Type) string {
	return b.typecast(t, false)
}

func (b *BaseJava) typecast(t *parser.Type, isplain bool) string {
	if t == nil {
		if isplain {
			return "void"
		} else {
			return "Void"
		}
	}

	var typemapping map[string]string

	if isplain {
		typemapping = plaintypemapping
	} else {
		typemapping = objecttypemapping
	}

	if t, ok := typemapping[t.Name]; ok {
		return t
	}

	switch t.Name {
	case langs.ThriftTypeList, langs.ThriftTypeSet:
		return fmt.Sprintf("ArrayList<%s>", b.ObjectTypecast(t.ValueType))
	case langs.ThriftTypeMap:
		return fmt.Sprintf("Map<%s, %s>", b.ObjectTypecast(t.KeyType), b.ObjectTypecast(t.ValueType))
	default:
		s := strings.Split(t.Name, ".")
		if len(s) == 1 {
			return s[0]
		} else if len(s) == 2 {
			pkg := ""
			for k, v := range b.t.Includes {
				if k == s[0] {
					for p, t := range *b.ts {
						if v == p {
							pkg = t.Namespaces[javaLang]
							break
						}
					}

					break
				}
			}

			if pkg == "" {
				return s[1]
			}
			return fmt.Sprintf("%s.%s", pkg, s[1])
		} else {
			return t.Name
		}
	}
}

func (b *BaseJava) AssembleParams(method *parser.Method) string {
	var buf bytes.Buffer

	for i, arg := range method.Arguments {
		if i != 0 {
			buf.WriteString(", ")
		}

		buf.WriteString(fmt.Sprintf("final %s %s", b.PlainTypecast(arg.Type), b.FilterVariableName(arg.Name)))
	}

	if len(method.Arguments) == 0 {
		buf.WriteString("")
	} else {
		buf.WriteString(", ")
	}

	buf.WriteString(fmt.Sprintf("final Listener<%s> listener", b.ObjectTypecast(method.ReturnType)))

	return buf.String()
}

func (b *BaseJava) GetInnerType(t *parser.Type) string {
	if t == nil {
		return "Void"
	}

	// map is ignored
	if t.Name == langs.ThriftTypeList || t.Name == langs.ThriftTypeSet {
		return b.GetInnerType(t.ValueType)
<<<<<<< fc3f70c15691348aa51baf9042726ad29ce4fea5
	}

	return b.ObjectTypecast(t)
}

// IsEnum checks whether a type is enum.
// it first checks in its own definition, than check from included files
func (b *BaseJava) IsEnum(t *parser.Type) bool {
	if t == nil {
		return false
	}

	names := strings.Split(t.Name, ".")

	if len(names) == 1 {
		for n := range b.t.Enums {
			if n == t.Name {
				return true
			}
		}

		return false
	}

	for path, thrift := range b.ts {
		if thrift == b.t {
			continue
		}

		filename := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
		if filename != names[0] {
			continue
		}

		for n := range thrift.Enums {
			if n == names[1] {
				return true
			}
		}
	}

	return false
}

type javaEnum struct {
	*BaseJava
	*parser.Enum
}

func (e *javaEnum) GenerateProperties() string {
	size := len(e.Enum.Values)
	i := 0

	var buf bytes.Buffer
	for _, v := range e.Enum.Values {
		buf.WriteString(fmt.Sprintf("\t@SerializedName(\"%d\")\n", v.Value))
		buf.WriteString(fmt.Sprintf("\t%s(%d)", strings.ToUpper(v.Name), v.Value))
		if i == size-1 {
			buf.WriteString(";\n")
		} else {
			buf.WriteString(",\n")
		}
		i++
	}
	return buf.String()
=======
	}

	return b.ObjectTypecast(t)
>>>>>>> change generated source files
}

type javaEnum struct {
	*BaseJava
	*parser.Enum
}

func (e *javaEnum) GenerateProperties() string {
	size := len(e.Enum.Values)
	i := 0

	var buf bytes.Buffer
	for _, v := range e.Enum.Values {
		buf.WriteString(fmt.Sprintf("\t@SerializedName(\"%d\")\n", v.Value))
		buf.WriteString(fmt.Sprintf("\t%s(%d)", strings.ToUpper(v.Name), v.Value))
		if i == size-1 {
			buf.WriteString(";\n")
		} else {
			buf.WriteString(",\n")
		}
		i++
	}
	return buf.String()
}

type javaStruct struct {
	*BaseJava
	*parser.Struct
}

func (this *javaStruct) HasKeyword() bool {
	for _, f := range this.Struct.Fields {
		if this.BaseJava.IsKeyword(f.Name) {
			return true
		}
	}
	return false
}

type javaService struct {
	*BaseJava
	*parser.Service
}

func generateAll(gen *JavaGen, output string, parsedThrift map[string]*parser.Thrift) {
	generateWithModel(gen, global.MODE_REST, filepath.Join(output, global.MODE_REST), parsedThrift)
	generateWithModel(gen, global.MODE_JSONRPC, filepath.Join(output, global.MODE_JSONRPC), parsedThrift)
}

func (g *JavaGen) Generate(output string, parsedThrift map[string]*parser.Thrift) {
	if global.Mode != "" {
		generateWithModel(g, global.Mode, output, parsedThrift)
	} else {
		generateAll(g, output, parsedThrift)
	}
}

func generateWithModel(gen *JavaGen, m string, output string, parsedThrift map[string]*parser.Thrift) {
	if m != global.MODE_REST && m != global.MODE_JSONRPC {
		log.Fatalf("mode '%s' is invalid", m)
	}

	gen.BaseGen.Init(javaLang, parsedThrift)

	if err := os.MkdirAll(output, 0755); err != nil {
		panic(fmt.Errorf("failed to create output directory %s", output))
	}

	// init templates
	enumTpl := initemplate(TPL_ENUM, "tmpl/java/enum.gojava")

	var structTpl *template.Template
	var serviceTpl *template.Template
	if m == global.MODE_REST {
		structTpl = initemplate(TPL_STRUCT, "tmpl/java/rest_struct.gojava")
		serviceTpl = initemplate(TPL_SERVICE, "tmpl/java/rest_service.gojava")
	} else if m == global.MODE_JSONRPC {
		structTpl = initemplate(TPL_STRUCT, "tmpl/java/jsonrpc_struct.gojava")
		serviceTpl = initemplate(TPL_SERVICE, "tmpl/java/jsonrpc_service.gojava")
	}

	// key is the absoule path of thrift file
	for _, t := range parsedThrift {
		// due to java's features,
		// we generate the struct and service in seperate template file

		ns := t.Namespaces["java"]

		if len(t.Enums) > 0 {
			log.Printf("## enums")

			for _, e := range t.Enums {
				name := e.Name + ".java"

				// fix java file path
				p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
				if err := os.MkdirAll(p, 0755); err != nil {
					panic(fmt.Errorf("failed to create output directory %s", p))
				}

				path := filepath.Join(p, name)

				base := BaseJava{Namespace: ns, t: t, ts: &parsedThrift}
				data := &javaEnum{BaseJava: &base, Enum: e}

				if err := outputfile(path, enumTpl, TPL_ENUM, data); err != nil {
					panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
				}

				log.Printf("%s", path)
			}
		}

		log.Printf("## structs")

		for _, s := range t.Structs {
			// filename is the struct name
			name := s.Name + javaExt

			// fix java file path
			p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
			if err := os.MkdirAll(p, 0755); err != nil {
				panic(fmt.Errorf("failed to create output directory %s", p))
			}

			path := filepath.Join(p, name)

			base := BaseJava{Namespace: ns, t: t, ts: parsedThrift}
			data := &javaStruct{BaseJava: &base, Struct: s}

			if err := outputfile(path, structTpl, TPL_STRUCT, data); err != nil {
				panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
			}

			log.Printf("%s", path)
		}

		log.Printf("## services")

		for _, s := range t.Services {
			// filename is the service name plus 'Service'
			name := s.Name + "Service" + javaExt

			// fix java file path
			p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
			if err := os.MkdirAll(p, 0755); err != nil {
				panic(fmt.Errorf("failed to create output directory %s", p))
			}

			path := filepath.Join(p, name)

			base := BaseJava{Namespace: ns, t: t, ts: parsedThrift}
			data := &javaService{BaseJava: &base, Service: s}

			if err := outputfile(path, serviceTpl, TPL_SERVICE, data); err != nil {
				panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
			}

			log.Printf("%s", path)
		}
	}
}

func initemplate(n string, path string) *template.Template {
	data, err := tmpl.Asset(path)
	if err != nil {
		panic(err)
	}

	tpl := template.New(n)

	// register functions
	funcMap := template.FuncMap{
		"ToUpper": strings.ToUpper,
	}
	tpl.Funcs(funcMap)

	_, err = tpl.Parse(string(data))
	if err != nil {
		panic(err)
	}

	return tpl
}

func outputfile(fp string, t *template.Template, tplname string, data interface{}) error {
	file, err := os.OpenFile(fp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	defer file.Close()

	return t.ExecuteTemplate(file, tplname, data)
}

func init() {
	langs.Langs[javaLang] = &JavaGen{}
}
<commit_msg>little refactor<commit_after>package java

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/ezbuy/tgen/global"
	"github.com/ezbuy/tgen/langs"
	"github.com/ezbuy/tgen/tmpl"
	"github.com/samuel/go-thrift/parser"
)

const (
	javaLang = "java"
	javaExt  = ".java"
)

const (
	JavaTypeshort  = "short"
	JavaTypeint    = "int"
	JavaTypelong   = "long"
	JavaTypebool   = "boolean"
	JavaTypebyte   = "byte"
	JavaTypedouble = "double"

	JavaTypeString = "String"

	JavaTypeShort  = "Short"
	JavaTypeInt    = "Integer"
	JavaTypeLong   = "Long"
	JavaTypeBool   = "Boolean"
	JavaTypeByte   = "Byte"
	JavaTypeDouble = "Double"

	// other types (such as array, map, etc.) are implemented in the method 'Typecast'
)

const (
	TPL_ENUM    = "tgen/java/enum"
	TPL_STRUCT  = "tgen/java/struct"
	TPL_SERVICE = "tgen/java/service"
)

var plaintypemapping = map[string]string{
	langs.ThriftTypeI16:    JavaTypeshort,
	langs.ThriftTypeI32:    JavaTypeint,
	langs.ThriftTypeI64:    JavaTypelong,
	langs.ThriftTypeString: JavaTypeString,
	langs.ThriftTypeByte:   JavaTypebyte,
	langs.ThriftTypeBool:   JavaTypebool,
	langs.ThriftTypeDouble: JavaTypedouble,
}

var objecttypemapping = map[string]string{
	langs.ThriftTypeI16:    JavaTypeShort,
	langs.ThriftTypeI32:    JavaTypeInt,
	langs.ThriftTypeI64:    JavaTypeLong,
	langs.ThriftTypeString: JavaTypeString,
	langs.ThriftTypeByte:   JavaTypeByte,
	langs.ThriftTypeBool:   JavaTypeBool,
	langs.ThriftTypeDouble: JavaTypeDouble,
}

type JavaGen struct {
	langs.BaseGen
}

type BaseJava struct {
	Namespace string
	t         *parser.Thrift
	ts        map[string]*parser.Thrift
}

func (b *BaseJava) FilterVariableName(n string) string {
	if b.IsKeyword(n) {
		return fmt.Sprintf("t%s%s", strings.ToUpper(n[:1]), n[1:])
	}
	return n
}

func (b *BaseJava) IsKeyword(n string) bool {
	switch n {
	case "package", "int", "short", "long", "byte", "boolean", "case", "switch", "if ", "for", "else",
		"goto", "Integer", "Short", "Long", "Byte", "Boolean", "class", "break", "try", "catch",
		"double", "Double", "do", "while", "final", "finally", "continue", "interface", "private",
		"public", "protected", "return", "this", "throw", "static", "super", "throws",
		"true", "false", "float", "volatile", "synchronized", "abstract", "default", "extends",
		"native", "new":
		return true
	}
	return false
}

func (b *BaseJava) PlainTypecast(t *parser.Type) string {
	return b.typecast(t, true)
}

func (b *BaseJava) ObjectTypecast(t *parser.Type) string {
	return b.typecast(t, false)
}

func (b *BaseJava) typecast(t *parser.Type, isplain bool) string {
	if t == nil {
		if isplain {
			return "void"
		} else {
			return "Void"
		}
	}

	var typemapping map[string]string

	if isplain {
		typemapping = plaintypemapping
	} else {
		typemapping = objecttypemapping
	}

	if t, ok := typemapping[t.Name]; ok {
		return t
	}

	switch t.Name {
	case langs.ThriftTypeList, langs.ThriftTypeSet:
		return fmt.Sprintf("ArrayList<%s>", b.ObjectTypecast(t.ValueType))
	case langs.ThriftTypeMap:
		return fmt.Sprintf("Map<%s, %s>", b.ObjectTypecast(t.KeyType), b.ObjectTypecast(t.ValueType))
	default:
		s := strings.Split(t.Name, ".")
		if len(s) == 1 {
			return s[0]
		} else if len(s) == 2 {
			pkg := ""
			for k, v := range b.t.Includes {
				if k == s[0] {
					for p, t := range *b.ts {
						if v == p {
							pkg = t.Namespaces[javaLang]
							break
						}
					}

					break
				}
			}

			if pkg == "" {
				return s[1]
			}
			return fmt.Sprintf("%s.%s", pkg, s[1])
		} else {
			return t.Name
		}
	}
}

func (b *BaseJava) AssembleParams(method *parser.Method) string {
	var buf bytes.Buffer

	for i, arg := range method.Arguments {
		if i != 0 {
			buf.WriteString(", ")
		}

		buf.WriteString(fmt.Sprintf("final %s %s", b.PlainTypecast(arg.Type), b.FilterVariableName(arg.Name)))
	}

	if len(method.Arguments) == 0 {
		buf.WriteString("")
	} else {
		buf.WriteString(", ")
	}

	buf.WriteString(fmt.Sprintf("final Listener<%s> listener", b.ObjectTypecast(method.ReturnType)))

	return buf.String()
}

func (b *BaseJava) GetInnerType(t *parser.Type) string {
	if t == nil {
		return "Void"
	}

	// map is ignored
	if t.Name == langs.ThriftTypeList || t.Name == langs.ThriftTypeSet {
		return b.GetInnerType(t.ValueType)
<<<<<<< fc3f70c15691348aa51baf9042726ad29ce4fea5
	}

	return b.ObjectTypecast(t)
}

// IsEnum checks whether a type is enum.
// it first checks in its own definition, than check from included files
func (b *BaseJava) IsEnum(t *parser.Type) bool {
	if t == nil {
		return false
	}

	names := strings.Split(t.Name, ".")

	if len(names) == 1 {
		for n := range b.t.Enums {
			if n == t.Name {
				return true
			}
		}

		return false
	}

	for path, thrift := range b.ts {
		if thrift == b.t {
			continue
		}

		filename := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
		if filename != names[0] {
			continue
		}

		for n := range thrift.Enums {
			if n == names[1] {
				return true
			}
		}
	}

	return false
}

type javaEnum struct {
	*BaseJava
	*parser.Enum
}

func (e *javaEnum) GenerateProperties() string {
	size := len(e.Enum.Values)
	i := 0

	var buf bytes.Buffer
	for _, v := range e.Enum.Values {
		buf.WriteString(fmt.Sprintf("\t@SerializedName(\"%d\")\n", v.Value))
		buf.WriteString(fmt.Sprintf("\t%s(%d)", strings.ToUpper(v.Name), v.Value))
		if i == size-1 {
			buf.WriteString(";\n")
		} else {
			buf.WriteString(",\n")
		}
		i++
	}
	return buf.String()
=======
	}

	return b.ObjectTypecast(t)
>>>>>>> change generated source files
}

type javaEnum struct {
	*BaseJava
	*parser.Enum
}

func (e *javaEnum) GenerateProperties() string {
	size := len(e.Enum.Values)
	i := 0

	var buf bytes.Buffer
	for _, v := range e.Enum.Values {
		buf.WriteString(fmt.Sprintf("\t@SerializedName(\"%d\")\n", v.Value))
		buf.WriteString(fmt.Sprintf("\t%s(%d)", strings.ToUpper(v.Name), v.Value))
		if i == size-1 {
			buf.WriteString(";\n")
		} else {
			buf.WriteString(",\n")
		}
		i++
	}
	return buf.String()
}

type javaStruct struct {
	*BaseJava
	*parser.Struct
}

func (this *javaStruct) HasKeyword() bool {
	for _, f := range this.Struct.Fields {
		if this.BaseJava.IsKeyword(f.Name) {
			return true
		}
	}
	return false
}

type javaService struct {
	*BaseJava
	*parser.Service
}

func generateAll(gen *JavaGen, output string, parsedThrift map[string]*parser.Thrift) {
	generateWithModel(gen, global.MODE_REST, filepath.Join(output, global.MODE_REST), parsedThrift)
	generateWithModel(gen, global.MODE_JSONRPC, filepath.Join(output, global.MODE_JSONRPC), parsedThrift)
}

func (g *JavaGen) Generate(output string, parsedThrift map[string]*parser.Thrift) {
	if global.Mode != "" {
		generateWithModel(g, global.Mode, output, parsedThrift)
	} else {
		generateAll(g, output, parsedThrift)
	}
}

func generateWithModel(gen *JavaGen, m string, output string, parsedThrift map[string]*parser.Thrift) {
	if m != global.MODE_REST && m != global.MODE_JSONRPC {
		log.Fatalf("mode '%s' is invalid", m)
	}

	gen.BaseGen.Init(javaLang, parsedThrift)

	if err := os.MkdirAll(output, 0755); err != nil {
		panic(fmt.Errorf("failed to create output directory %s", output))
	}

	// init templates
	enumTpl := initemplate(TPL_ENUM, "tmpl/java/enum.gojava")

	var structTpl *template.Template
	var serviceTpl *template.Template
	if m == global.MODE_REST {
		structTpl = initemplate(TPL_STRUCT, "tmpl/java/rest_struct.gojava")
		serviceTpl = initemplate(TPL_SERVICE, "tmpl/java/rest_service.gojava")
	} else if m == global.MODE_JSONRPC {
		structTpl = initemplate(TPL_STRUCT, "tmpl/java/jsonrpc_struct.gojava")
		serviceTpl = initemplate(TPL_SERVICE, "tmpl/java/jsonrpc_service.gojava")
	}

	// key is the absoule path of thrift file
	for _, t := range parsedThrift {
		// due to java's features,
		// we generate the struct and service in seperate template file

		ns := t.Namespaces[javaLang]

		if len(t.Enums) > 0 {
			log.Printf("## enums")

			for _, e := range t.Enums {
				name := e.Name + javaExt

				// fix java file path
				p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
				if err := os.MkdirAll(p, 0755); err != nil {
					panic(fmt.Errorf("failed to create output directory %s", p))
				}

				path := filepath.Join(p, name)

				base := BaseJava{Namespace: ns, t: t, ts: &parsedThrift}
				data := &javaEnum{BaseJava: &base, Enum: e}

				if err := outputfile(path, enumTpl, TPL_ENUM, data); err != nil {
					panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
				}

				log.Printf("%s", path)
			}
		}

		log.Printf("## structs")

		for _, s := range t.Structs {
			// filename is the struct name
			name := s.Name + javaExt

			// fix java file path
			p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
			if err := os.MkdirAll(p, 0755); err != nil {
				panic(fmt.Errorf("failed to create output directory %s", p))
			}

			path := filepath.Join(p, name)

			base := BaseJava{Namespace: ns, t: t, ts: parsedThrift}
			data := &javaStruct{BaseJava: &base, Struct: s}

			if err := outputfile(path, structTpl, TPL_STRUCT, data); err != nil {
				panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
			}

			log.Printf("%s", path)
		}

		log.Printf("## services")

		for _, s := range t.Services {
			// filename is the service name plus 'Service'
			name := s.Name + "Service" + javaExt

			// fix java file path
			p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
			if err := os.MkdirAll(p, 0755); err != nil {
				panic(fmt.Errorf("failed to create output directory %s", p))
			}

			path := filepath.Join(p, name)

			base := BaseJava{Namespace: ns, t: t, ts: parsedThrift}
			data := &javaService{BaseJava: &base, Service: s}

			if err := outputfile(path, serviceTpl, TPL_SERVICE, data); err != nil {
				panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
			}

			log.Printf("%s", path)
		}
	}
}

func initemplate(n string, path string) *template.Template {
	data, err := tmpl.Asset(path)
	if err != nil {
		panic(err)
	}

	tpl := template.New(n)

	// register functions
	funcMap := template.FuncMap{
		"ToUpper": strings.ToUpper,
	}
	tpl.Funcs(funcMap)

	_, err = tpl.Parse(string(data))
	if err != nil {
		panic(err)
	}

	return tpl
}

func outputfile(fp string, t *template.Template, tplname string, data interface{}) error {
	file, err := os.OpenFile(fp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	defer file.Close()

	return t.ExecuteTemplate(file, tplname, data)
}

func init() {
	langs.Langs[javaLang] = &JavaGen{}
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"github.com/docopt/docopt-go"
	_ "github.com/mattn/go-sqlite3"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
)

const (
	dbRelativePath = "Library/Application Support/Dock/desktoppicture.db"
	dbStatements   = `
		delete from data;
		delete from preferences;
		insert into data values (?);
		insert into data values (?);
		insert into preferences select 1, 1, ROWID from pictures;
		insert into preferences select 2, 2, ROWID from pictures;
	`
	usage = `Sets wallpaper to <wallpaper>. Fills the screen by default.

Usage:
  %[1]s [-f | -s | -c | -t] <wallpaper>
  %[1]s -h | -v

Options:
  -f --fit      Fit wallpaper to screen.
  -s --stretch  Stretch wallpaper to fill screen.
  -c --center   Center wallpaper.
  -t --tile     Tile wallpaper.
  -h --help     Show this help message.
  -v --version  Show version information.

`
	version = "%s version 0.1"
)

var positions = [...]string{"--fit", "--stretch", "--center", "--tile"}

func main() {
	log.SetFlags(0)

	args, err := docopt.Parse(fmt.Sprintf(usage, os.Args[0]), nil, true, fmt.Sprintf(version, os.Args[0]), false)
	if err != nil {
		log.Fatalf("cannot parse arguments (%s)", err)
	}

	wallpaper := args["<wallpaper>"]

	position := 5
	for _, p := range positions {
		if args[p].(bool) {
			break
		}
		position--
	}

	home, err := homeDir()
	if err != nil {
		log.Fatalf("cannot open database", err)
	}
	dbPath := filepath.Join(home, dbRelativePath)

	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		log.Fatalf("cannot open database (%s)", err)
	}
	defer db.Close()

	if _, err := db.Exec(dbStatements, wallpaper, position); err != nil {
		log.Fatalf("error updating database (%s)", err)
	}

	if err := exec.Command("killall", "Dock").Run(); err != nil {
		log.Println("error killing Dock, wallpaper will be applied on your next login")
	}
}

func homeDir() (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", fmt.Errorf("unknown current user")
	}
	if user.HomeDir == "" {
		return "", fmt.Errorf("unknown home directory")
	}
	return user.HomeDir, nil
}
<commit_msg>improve error message<commit_after>package main

import (
	"database/sql"
	"fmt"
	"github.com/docopt/docopt-go"
	_ "github.com/mattn/go-sqlite3"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
)

const (
	dbRelativePath = "Library/Application Support/Dock/desktoppicture.db"
	dbStatements   = `
		delete from data;
		delete from preferences;
		insert into data values (?);
		insert into data values (?);
		insert into preferences select 1, 1, ROWID from pictures;
		insert into preferences select 2, 2, ROWID from pictures;
	`
	usage = `Sets wallpaper to <wallpaper>. Fills the screen by default.

Usage:
  %[1]s [-f | -s | -c | -t] <wallpaper>
  %[1]s -h | -v

Options:
  -f --fit      Fit wallpaper to screen.
  -s --stretch  Stretch wallpaper to fill screen.
  -c --center   Center wallpaper.
  -t --tile     Tile wallpaper.
  -h --help     Show this help message.
  -v --version  Show version information.

`
	version = "%s version 0.1"
)

var positions = [...]string{"--fit", "--stretch", "--center", "--tile"}

func main() {
	log.SetFlags(0)

	args, err := docopt.Parse(fmt.Sprintf(usage, os.Args[0]), nil, true, fmt.Sprintf(version, os.Args[0]), false)
	if err != nil {
		log.Fatalf("cannot parse arguments (%s)", err)
	}

	wallpaper := args["<wallpaper>"]

	position := 5
	for _, p := range positions {
		if args[p].(bool) {
			break
		}
		position--
	}

	home, err := homeDir()
	if err != nil {
		log.Fatalf("cannot open database", err)
	}
	dbPath := filepath.Join(home, dbRelativePath)

	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		log.Fatalf("cannot open database (%s)", err)
	}
	defer db.Close()

	if _, err := db.Exec(dbStatements, wallpaper, position); err != nil {
		log.Fatalf("error updating database (%s)", err)
	}

	if err := exec.Command("killall", "Dock").Run(); err != nil {
		log.Println("error applying wallpaper, it will be applied on your next login")
	}
}

func homeDir() (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", fmt.Errorf("unknown current user")
	}
	if user.HomeDir == "" {
		return "", fmt.Errorf("unknown home directory")
	}
	return user.HomeDir, nil
}
<|endoftext|>


<commit_before>package disgo

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/slf4go/logger"
)

type shard struct {
	// Parent session
	session *Session

	// Websocket information used for identification and reconnection
	webSocket *websocket.Conn
	shard     int
	sessionID string
	sequence  int
	heartbeat int

	// Mutex locks, reconnect to make sure there is only 1 process reconnecting and concurrent read/write accesses on the socket
	reconnectLock       sync.Mutex
	readLock, writeLock sync.Mutex

	// Channels to pass around messages
	closeMessage chan int
	stopListen   chan bool
	stopRead     chan bool
}

func connectShard(session *Session, shardNum int) (*shard, error) {
	conn, _, err := websocket.DefaultDialer.Dial(session.wsUrl, http.Header{})
	if err != nil {
		return nil, err
	}

	s := &shard{
		session:      session,
		webSocket:    conn,
		shard:        shardNum,
		closeMessage: make(chan int, 1),
		stopListen:   make(chan bool),
		stopRead:     make(chan bool),
	}

	if err = s.handshake(); err != nil {
		return nil, err
	}

	return s, nil
}

func (s *shard) handshake() error {
	s.webSocket.SetCloseHandler(s.onClose)

	helloFrame, err := s.readFrame()
	if err != nil {
		return err
	} else if helloFrame.Op != opHello {
		return errors.New("First frame sent from the Discord Gateway is not hello")
	}

	hello := helloPayload{}
	if err = json.Unmarshal(helloFrame.Data, &hello); err != nil {
		return err
	}

	logger.Debugf("Connected to Discord servers: %s", strings.Join(hello.Servers, ", "))
	logger.Debugf("Setting up a heartbeat interval of %d ms", hello.HeartbeatInterval)
	s.heartbeat = hello.HeartbeatInterval

	if err = s.identify(); err != nil {
		return err
	}

	go s.mainLoop()
	return nil
}

func (s *shard) identify() error {
	if s.sessionID != "" {
		logger.Debugf("Resuming connection starting at sequence %d.", s.sequence)
		s.sendFrame(&gatewayFrame{opResume, resumePayload{
			Token:     s.session.token,
			SessionID: s.sessionID,
			Sequence:  s.sequence,
		}})
	} else {
		logger.Debugf("Identifying to websocket")
		s.sendFrame(&gatewayFrame{opIdentify, identifyPayload{
			Token:          s.session.token,
			Compress:       true,
			LargeThreshold: 250,
			Shard:          [2]int{s.shard, cap(s.session.shards)},
			Properties: propertiesPayload{
				OS:      runtime.GOOS,
				Browser: "DisGo",
				Device:  "DisGo",
			},
		}})
	}

	for {
		frame, err := s.readFrame()
		if err != nil {
			return err
		}

		if frame.Op == opDispatch {
			switch frame.EventName {
			case "READY":
				ready := ReadyEvent{}
				if err := json.Unmarshal(frame.Data, &ready); err != nil {
					return err
				}

				s.sessionID = ready.SessionID

				fallthrough
			case "RESUMED":
				s.session.dispatchEvent(frame)
				return nil // Break out of the loop, we have what we want
			default:
				s.session.dispatchEvent(frame) // Nope, resume action, let's wait for more frames
			}
		} else if frame.Op == opInvalidSession {
			s.sessionID = "" // Invalidate session and retry
			s.identify()
			return nil
		} else {
			return fmt.Errorf("Unexpected opCode received from Discord: %d", frame.Op)
		}
	}
}

func (s *shard) mainLoop() {
	logger.Debugf("Starting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	heartbeat := time.NewTicker(time.Duration(s.heartbeat) * time.Millisecond)
	defer heartbeat.Stop()
	sentHeartBeat := false

	reader := make(chan *receivedFrame, 1)
	go s.readWebSocket(reader)

	for {
		select {
		case <-heartbeat.C:
			if !sentHeartBeat {
				go s.sendFrame(&gatewayFrame{opHeartbeat, s.sequence})
				sentHeartBeat = true
			} else {
				go s.disconnect(websocket.CloseAbnormalClosure, "Did not respond to previous heartbeat")
			}
		case <-s.stopListen:
			return
		case frame := <-reader:
			switch frame.Op {
			case opHeartbeat:
				s.sendFrame(&gatewayFrame{Op: opHeartbeatAck})
			case opHeartbeatAck:
				sentHeartBeat = false
			case opReconnect:
				s.disconnect(websocket.CloseNormalClosure, "op Reconnect")
			case opDispatch:
				s.session.dispatchEvent(frame)
			default:
				logger.Errorf("Unknown opCode received: %d", frame.Op)
			}
		}
	}
}

func (s *shard) sendFrame(frame *gatewayFrame) {
	// Resume will only be sent within the reconnect lock, we don't relock to prevent a deadlock
	if frame.Op != opResume {
		s.reconnectLock.Lock()
		defer s.reconnectLock.Unlock()
	}

	s.writeLock.Lock()
	defer s.writeLock.Unlock()

	logger.Debugf("Sending frame with opCode: %d", frame.Op)
	s.webSocket.WriteJSON(frame)
}

func (s *shard) readWebSocket(reader chan *receivedFrame) {
	logger.Debugf("Starting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	for {
		select {
		case <-s.stopRead:
			return
		default:
			frame, err := s.readFrame()
			if err != nil {
				if !s.session.isShuttingDown() {
					if closeErr, isClose := err.(*websocket.CloseError); isClose {
						s.closeMessage <- closeErr.Code
					}
					logger.ErrorE(err)
					go s.disconnect(websocket.CloseAbnormalClosure, err.Error())
				}
				<-s.stopRead // Wait for the connection to be closed
				return
			}

			reader <- frame
		}
	}

}

func (s *shard) readFrame() (*receivedFrame, error) {
	s.readLock.Lock()
	defer s.readLock.Unlock()

	logger.Tracef("Shard.readFrame() called")
	msgType, msg, err := s.webSocket.ReadMessage()
	if err != nil {
		return nil, err
	}

	var reader io.Reader
	reader = bytes.NewBuffer(msg)

	if msgType == websocket.BinaryMessage {
		zReader, err := zlib.NewReader(reader)
		if err != nil {
			return nil, err
		}

		defer zReader.Close()

		reader = zReader
	}

	frame := receivedFrame{Sequence: -1}
	json.NewDecoder(reader).Decode(&frame)

	logger.Debugf("Received frame with opCode: %d", frame.Op)

	if frame.Sequence != -1 {
		logger.Tracef("Last sequence received set to %d.", frame.Sequence)
		s.sequence = frame.Sequence
	}
	return &frame, nil
}

func (s *shard) reconnect(recursive bool) {
	if !recursive {
		s.reconnectLock.Lock()
	}

	if !s.session.isShuttingDown() {
		logger.Noticef("Reconnecting shard [%d/%d]", s.shard+1, cap(s.session.shards))
		conn, _, err := websocket.DefaultDialer.Dial(s.session.wsUrl, http.Header{})

		if err == nil {
			s.webSocket = conn
			err = s.handshake()

			if err != nil {
				conn.Close()
				s.webSocket = nil
			}
		}

		if err != nil {
			logger.ErrorE(err)
			time.Sleep(1 * time.Second)
			go s.reconnect(true)
		} else {
			s.reconnectLock.Unlock()
		}
	}
}

func (s *shard) onClose(code int, text string) error {
	logger.Debugf("Received Close Frame from Discord. Code: %d. Text: %s", code, text)
	s.closeMessage <- code
	go s.reconnect(false)
	return nil
}

func (s *shard) disconnect(code int, text string) {
	if s.webSocket == nil {
		logger.Warnf("Shard.disconnect() called on a connection that is already gone, ignoring")
		return
	}

	logger.Tracef("Shard.disconnect() called")
	s.stopListen <- true

	s.writeLock.Lock()
	err := s.webSocket.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, text))
	s.writeLock.Unlock()
	if err != nil {
		logger.ErrorE(err)
	}

	select {
	case code := <-s.closeMessage:
		logger.Debugf("Discord connection closed with code %d", code)
	case <-time.After(1 * time.Second):
		logger.Warn("Discord did not reply to the close message, force-closing connection.")
	}

	s.webSocket.Close()
	s.stopRead <- true

	go s.reconnect(false)
}
<commit_msg>Websocket is now never set to nil to avoid sigsegv-panics<commit_after>package disgo

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/slf4go/logger"
)

type shard struct {
	// Parent session
	session *Session

	// Websocket information used for identification and reconnection
	webSocket *websocket.Conn
	shard     int
	sessionID string
	sequence  int
	heartbeat int

	// Mutex locks, reconnect to make sure there is only 1 process reconnecting and concurrent read/write accesses on the socket
	reconnectLock       sync.Mutex
	readLock, writeLock sync.Mutex

	// Channels to pass around messages
	closeMessage chan int
	stopListen   chan bool
	stopRead     chan bool
}

func connectShard(session *Session, shardNum int) (*shard, error) {
	conn, _, err := websocket.DefaultDialer.Dial(session.wsUrl, http.Header{})
	if err != nil {
		return nil, err
	}

	s := &shard{
		session:      session,
		webSocket:    conn,
		shard:        shardNum,
		closeMessage: make(chan int, 1),
		stopListen:   make(chan bool),
		stopRead:     make(chan bool),
	}

	if err = s.handshake(); err != nil {
		return nil, err
	}

	return s, nil
}

func (s *shard) handshake() error {
	s.webSocket.SetCloseHandler(s.onClose)

	helloFrame, err := s.readFrame()
	if err != nil {
		return err
	} else if helloFrame.Op != opHello {
		return errors.New("First frame sent from the Discord Gateway is not hello")
	}

	hello := helloPayload{}
	if err = json.Unmarshal(helloFrame.Data, &hello); err != nil {
		return err
	}

	logger.Debugf("Connected to Discord servers: %s", strings.Join(hello.Servers, ", "))
	logger.Debugf("Setting up a heartbeat interval of %d ms", hello.HeartbeatInterval)
	s.heartbeat = hello.HeartbeatInterval

	if err = s.identify(); err != nil {
		return err
	}

	go s.mainLoop()
	return nil
}

func (s *shard) identify() error {
	if s.sessionID != "" {
		logger.Debugf("Resuming connection starting at sequence %d.", s.sequence)
		s.sendFrame(&gatewayFrame{opResume, resumePayload{
			Token:     s.session.token,
			SessionID: s.sessionID,
			Sequence:  s.sequence,
		}})
	} else {
		logger.Debugf("Identifying to websocket")
		s.sendFrame(&gatewayFrame{opIdentify, identifyPayload{
			Token:          s.session.token,
			Compress:       true,
			LargeThreshold: 250,
			Shard:          [2]int{s.shard, cap(s.session.shards)},
			Properties: propertiesPayload{
				OS:      runtime.GOOS,
				Browser: "DisGo",
				Device:  "DisGo",
			},
		}})
	}

	for {
		frame, err := s.readFrame()
		if err != nil {
			return err
		}

		if frame.Op == opDispatch {
			switch frame.EventName {
			case "READY":
				ready := ReadyEvent{}
				if err := json.Unmarshal(frame.Data, &ready); err != nil {
					return err
				}

				s.sessionID = ready.SessionID

				fallthrough
			case "RESUMED":
				s.session.dispatchEvent(frame)
				return nil // Break out of the loop, we have what we want
			default:
				s.session.dispatchEvent(frame) // Nope, resume action, let's wait for more frames
			}
		} else if frame.Op == opInvalidSession {
			s.sessionID = "" // Invalidate session and retry
			s.identify()
			return nil
		} else {
			return fmt.Errorf("Unexpected opCode received from Discord: %d", frame.Op)
		}
	}
}

func (s *shard) mainLoop() {
	logger.Debugf("Starting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	heartbeat := time.NewTicker(time.Duration(s.heartbeat) * time.Millisecond)
	defer heartbeat.Stop()
	sentHeartBeat := false

	reader := make(chan *receivedFrame, 1)
	go s.readWebSocket(reader)

	for {
		select {
		case <-heartbeat.C:
			if !sentHeartBeat {
				go s.sendFrame(&gatewayFrame{opHeartbeat, s.sequence})
				sentHeartBeat = true
			} else {
				go s.disconnect(websocket.CloseAbnormalClosure, "Did not respond to previous heartbeat")
			}
		case <-s.stopListen:
			return
		case frame := <-reader:
			switch frame.Op {
			case opHeartbeat:
				s.sendFrame(&gatewayFrame{Op: opHeartbeatAck})
			case opHeartbeatAck:
				sentHeartBeat = false
			case opReconnect:
				s.disconnect(websocket.CloseNormalClosure, "op Reconnect")
			case opDispatch:
				s.session.dispatchEvent(frame)
			default:
				logger.Errorf("Unexpected opCode received: %d", frame.Op)
			}
		}
	}
}

func (s *shard) sendFrame(frame *gatewayFrame) {
	// Resume will only be sent within the reconnect lock, we don't relock to prevent a deadlock
	if frame.Op != opResume {
		s.reconnectLock.Lock()
		defer s.reconnectLock.Unlock()
	}

	s.writeLock.Lock()
	defer s.writeLock.Unlock()

	logger.Debugf("Sending frame with opCode: %d", frame.Op)
	s.webSocket.WriteJSON(frame)
}

func (s *shard) readWebSocket(reader chan *receivedFrame) {
	logger.Debugf("Starting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	for {
		select {
		case <-s.stopRead:
			return
		default:
			frame, err := s.readFrame()
			if err != nil {
				if !s.session.isShuttingDown() {
					if closeErr, isClose := err.(*websocket.CloseError); isClose {
						s.closeMessage <- closeErr.Code
					}
					logger.ErrorE(err)
					go s.disconnect(websocket.CloseAbnormalClosure, err.Error())
				}
				<-s.stopRead // Wait for the connection to be closed
				return
			}

			reader <- frame
		}
	}

}

func (s *shard) readFrame() (*receivedFrame, error) {
	s.readLock.Lock()
	defer s.readLock.Unlock()

	logger.Tracef("Shard.readFrame() called")
	msgType, msg, err := s.webSocket.ReadMessage()
	if err != nil {
		return nil, err
	}

	var reader io.Reader
	reader = bytes.NewBuffer(msg)

	if msgType == websocket.BinaryMessage {
		zReader, err := zlib.NewReader(reader)
		if err != nil {
			return nil, err
		}

		defer zReader.Close()

		reader = zReader
	}

	frame := receivedFrame{Sequence: -1}
	json.NewDecoder(reader).Decode(&frame)

	logger.Debugf("Received frame with opCode: %d", frame.Op)

	if frame.Sequence != -1 {
		logger.Tracef("Last sequence received set to %d.", frame.Sequence)
		s.sequence = frame.Sequence
	}
	return &frame, nil
}

func (s *shard) reconnect(recursive bool) {
	if !recursive {
		s.reconnectLock.Lock()
	}

	if !s.session.isShuttingDown() {
		logger.Noticef("Reconnecting shard [%d/%d]", s.shard+1, cap(s.session.shards))
		conn, _, err := websocket.DefaultDialer.Dial(s.session.wsUrl, http.Header{})

		if err == nil {
			s.webSocket = conn
			err = s.handshake()

			if err != nil {
				conn.Close()
			}
		}

		if err != nil {
			logger.ErrorE(err)
			time.Sleep(1 * time.Second)
			go s.reconnect(true)
		} else {
			s.reconnectLock.Unlock()
		}
	}
}

func (s *shard) onClose(code int, text string) error {
	logger.Debugf("Received Close Frame from Discord. Code: %d. Text: %s", code, text)
	s.closeMessage <- code
	go s.reconnect(false)
	return nil
}

func (s *shard) disconnect(code int, text string) {
	logger.Tracef("Shard.disconnect() called")
	s.stopListen <- true

	s.writeLock.Lock()
	err := s.webSocket.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, text))
	s.writeLock.Unlock()
	if err != nil {
		logger.ErrorE(err)
	}

	select {
	case code := <-s.closeMessage:
		logger.Debugf("Discord connection closed with code %d", code)
	case <-time.After(1 * time.Second):
		logger.Warn("Discord did not reply to the close message, force-closing connection.")
	}

	s.webSocket.Close()
	s.stopRead <- true

	go s.reconnect(false)
}
<|endoftext|>


<commit_before>package disgo

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/slf4go/logger"
)

type shard struct {
	session *Session

	webSocket *websocket.Conn
	lock      sync.Mutex
	shard     int

	sessionID string
	sequence  int
	heartbeat int

	closeMessage chan int
	stopListen   chan bool
	stopRead     chan bool
}

func connectShard(session *Session, shardNum int) (*shard, error) {
	conn, _, err := websocket.DefaultDialer.Dial(session.wsUrl, http.Header{})
	if err != nil {
		return nil, err
	}

	s := &shard{
		session:      session,
		webSocket:    conn,
		shard:        shardNum,
		closeMessage: make(chan int, 1),
		stopListen:   make(chan bool),
		stopRead:     make(chan bool),
	}

	if err = s.handshake(); err != nil {
		return nil, err
	}

	return s, nil
}

func (s *shard) handshake() error {
	s.webSocket.SetCloseHandler(s.onClose)

	helloFrame, err := s.readFrame()
	if err != nil {
		return err
	} else if helloFrame.Op != opHello {
		return errors.New("First frame sent from the Discord Gateway is not hello")
	}

	hello := helloPayload{}
	if err = json.Unmarshal(helloFrame.Data, &hello); err != nil {
		return err
	}

	logger.Debugf("Connected to Discord servers: %s", strings.Join(hello.Servers, ", "))
	logger.Debugf("Setting up a heartbeat interval of %d ms", hello.HeartbeatInterval)
	s.heartbeat = hello.HeartbeatInterval

	if err = s.identify(); err != nil {
		return err
	}

	go s.mainLoop()
	return nil
}

func (s *shard) identify() error {
	if s.sessionID != "" {
		logger.Debugf("Resuming connection starting at sequence %d.", s.sequence)
		s.sendFrame(&gatewayFrame{opResume, resumePayload{
			Token:     s.session.token,
			SessionID: s.sessionID,
			Sequence:  s.sequence,
		}})
	} else {
		logger.Debugf("Identifying to websocket")
		s.sendFrame(&gatewayFrame{opIdentify, identifyPayload{
			Token:          s.session.token,
			Compress:       true,
			LargeThreshold: 250,
			Shard:          [2]int{s.shard, cap(s.session.shards)},
			Properties: propertiesPayload{
				OS:      runtime.GOOS,
				Browser: "DisGo",
				Device:  "DisGo",
			},
		}})
	}

	frame, err := s.readFrame()
	if err != nil {
		return err
	}

	if frame.Op == opDispatch {
		if frame.EventName == "READY" {
			ready := ReadyEvent{}
			if err := json.Unmarshal(frame.Data, &ready); err != nil {
				return err
			}

			s.sessionID = ready.SessionID
			s.session.dispatchEvent(frame)
			return nil
		} else if frame.EventName == "RESUMED" {
			s.session.dispatchEvent(frame)
			return nil
		} else {
			return fmt.Errorf("Discord sent event of type '%s', expected 'READY' or 'RESUMED'", frame.EventName)
		}
	} else if frame.Op == opInvalidSession {
		s.sessionID = "" // Invalidate session and retry
		s.identify()
		return nil
	} else {
		return fmt.Errorf("Unexpected opCode received from Discord: %d", frame.Op)
	}
}

func (s *shard) mainLoop() {
	logger.Debugf("Starting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	heartbeat := time.NewTicker(time.Duration(s.heartbeat) * time.Millisecond)
	defer heartbeat.Stop()
	sentHeartBeat := false

	reader := make(chan *receivedFrame, 1)
	go s.readWebSocket(reader)

	for {
		select {
		case <-heartbeat.C:
			if !sentHeartBeat {
				s.sendFrame(&gatewayFrame{opHeartbeat, s.sequence})
				sentHeartBeat = true
			} else {
				s.disconnect(websocket.CloseAbnormalClosure, "Did not respond to previous ping")
			}
		case <-s.stopListen:
			return
		case frame := <-reader:
			switch frame.Op {
			case opHeartbeat:
				s.sendFrame(&gatewayFrame{Op: opHeartbeatAck})
			case opHeartbeatAck:
				sentHeartBeat = false
			case opReconnect:
				s.disconnect(websocket.CloseNormalClosure, "op Reconnect")
			case opDispatch:
				s.session.dispatchEvent(frame)
			default:
				logger.Errorf("Unknown opCode received: %d", frame.Op)
			}
		}
	}
}

func (s *shard) sendFrame(frame *gatewayFrame) {
	// Resume will only be sent within the reconnect lock, we don't relock to prevent a deadlock
	if frame.Op != opResume {
		s.lock.Lock()
		defer s.lock.Unlock()
	}

	logger.Debugf("Sending frame with opCode: %d", frame.Op)
	s.webSocket.WriteJSON(frame)
}

func (s *shard) readWebSocket(reader chan *receivedFrame) {
	logger.Debugf("Starting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	for {
		select {
		case <-s.stopRead:
			return
		default:
			frame, err := s.readFrame()
			if err != nil {
				if !s.session.isShuttingDown() {
					if closeErr, isClose := err.(*websocket.CloseError); isClose {
						s.closeMessage <- closeErr.Code
					}
					logger.ErrorE(err)
					go s.disconnect(websocket.CloseAbnormalClosure, err.Error())
				}
				<-s.stopRead // Wait for the connection to be closed
				return
			}

			reader <- frame
		}
	}

}

func (s *shard) readFrame() (*receivedFrame, error) {
	logger.Tracef("Shard.readFrame() called")
	msgType, msg, err := s.webSocket.ReadMessage()
	if err != nil {
		return nil, err
	}

	var reader io.Reader
	reader = bytes.NewBuffer(msg)

	if msgType == websocket.BinaryMessage {
		zReader, err := zlib.NewReader(reader)
		if err != nil {
			return nil, err
		}

		defer zReader.Close()

		reader = zReader
	}

	frame := receivedFrame{Sequence: -1}
	json.NewDecoder(reader).Decode(&frame)

	logger.Debugf("Received frame with opCode: %d", frame.Op)

	if frame.Sequence != -1 {
		logger.Tracef("Last sequence received set to %d.", frame.Sequence)
		s.sequence = frame.Sequence
	}
	return &frame, nil
}

func (s *shard) reconnect() {
	s.lock.Lock()
	defer s.lock.Unlock()

	if !s.session.isShuttingDown() {
		logger.Noticef("Reconnecting shard [%d/%d]", s.shard+1, cap(s.session.shards))
		conn, _, err := websocket.DefaultDialer.Dial(s.session.wsUrl, http.Header{})

		if err == nil {
			s.webSocket = conn
			err = s.handshake()

			if err != nil {
				conn.Close()
				s.webSocket = nil
			}
		}

		if err != nil {
			logger.ErrorE(err)
			time.Sleep(1 * time.Second)
			go s.reconnect()
		}
	}
}

func (s *shard) onClose(code int, text string) error {
	logger.Debugf("Received Close Frame from Discord. Code: %d. Text: %s", code, text)
	s.closeMessage <- code
	s.reconnect()
	return nil
}

func (s *shard) disconnect(code int, text string) {
	logger.Tracef("Shard.disconnect() called")
	s.stopListen <- true

	s.lock.Lock()
	err := s.webSocket.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, text))
	s.lock.Unlock()
	if err != nil {
		logger.ErrorE(err)
	}

	select {
	case code := <-s.closeMessage:
		logger.Debugf("Discord connection closed with code %d", code)
	case <-time.After(1 * time.Second):
		logger.Warn("Discord did not reply to the close message, force-closing connection.")
	}

	s.webSocket.Close()
	s.stopRead <- true

	go s.reconnect()
}
<commit_msg>Reconnect now holds the lock inbetween attempts<commit_after>package disgo

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/slf4go/logger"
)

type shard struct {
	session *Session

	webSocket *websocket.Conn
	lock      sync.Mutex
	shard     int

	sessionID string
	sequence  int
	heartbeat int

	closeMessage chan int
	stopListen   chan bool
	stopRead     chan bool
}

func connectShard(session *Session, shardNum int) (*shard, error) {
	conn, _, err := websocket.DefaultDialer.Dial(session.wsUrl, http.Header{})
	if err != nil {
		return nil, err
	}

	s := &shard{
		session:      session,
		webSocket:    conn,
		shard:        shardNum,
		closeMessage: make(chan int, 1),
		stopListen:   make(chan bool),
		stopRead:     make(chan bool),
	}

	if err = s.handshake(); err != nil {
		return nil, err
	}

	return s, nil
}

func (s *shard) handshake() error {
	s.webSocket.SetCloseHandler(s.onClose)

	helloFrame, err := s.readFrame()
	if err != nil {
		return err
	} else if helloFrame.Op != opHello {
		return errors.New("First frame sent from the Discord Gateway is not hello")
	}

	hello := helloPayload{}
	if err = json.Unmarshal(helloFrame.Data, &hello); err != nil {
		return err
	}

	logger.Debugf("Connected to Discord servers: %s", strings.Join(hello.Servers, ", "))
	logger.Debugf("Setting up a heartbeat interval of %d ms", hello.HeartbeatInterval)
	s.heartbeat = hello.HeartbeatInterval

	if err = s.identify(); err != nil {
		return err
	}

	go s.mainLoop()
	return nil
}

func (s *shard) identify() error {
	if s.sessionID != "" {
		logger.Debugf("Resuming connection starting at sequence %d.", s.sequence)
		s.sendFrame(&gatewayFrame{opResume, resumePayload{
			Token:     s.session.token,
			SessionID: s.sessionID,
			Sequence:  s.sequence,
		}})
	} else {
		logger.Debugf("Identifying to websocket")
		s.sendFrame(&gatewayFrame{opIdentify, identifyPayload{
			Token:          s.session.token,
			Compress:       true,
			LargeThreshold: 250,
			Shard:          [2]int{s.shard, cap(s.session.shards)},
			Properties: propertiesPayload{
				OS:      runtime.GOOS,
				Browser: "DisGo",
				Device:  "DisGo",
			},
		}})
	}

	frame, err := s.readFrame()
	if err != nil {
		return err
	}

	if frame.Op == opDispatch {
		if frame.EventName == "READY" {
			ready := ReadyEvent{}
			if err := json.Unmarshal(frame.Data, &ready); err != nil {
				return err
			}

			s.sessionID = ready.SessionID
			s.session.dispatchEvent(frame)
			return nil
		} else if frame.EventName == "RESUMED" {
			s.session.dispatchEvent(frame)
			return nil
		} else {
			return fmt.Errorf("Discord sent event of type '%s', expected 'READY' or 'RESUMED'", frame.EventName)
		}
	} else if frame.Op == opInvalidSession {
		s.sessionID = "" // Invalidate session and retry
		s.identify()
		return nil
	} else {
		return fmt.Errorf("Unexpected opCode received from Discord: %d", frame.Op)
	}
}

func (s *shard) mainLoop() {
	logger.Debugf("Starting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	heartbeat := time.NewTicker(time.Duration(s.heartbeat) * time.Millisecond)
	defer heartbeat.Stop()
	sentHeartBeat := false

	reader := make(chan *receivedFrame, 1)
	go s.readWebSocket(reader)

	for {
		select {
		case <-heartbeat.C:
			if !sentHeartBeat {
				s.sendFrame(&gatewayFrame{opHeartbeat, s.sequence})
				sentHeartBeat = true
			} else {
				s.disconnect(websocket.CloseAbnormalClosure, "Did not respond to previous ping")
			}
		case <-s.stopListen:
			return
		case frame := <-reader:
			switch frame.Op {
			case opHeartbeat:
				s.sendFrame(&gatewayFrame{Op: opHeartbeatAck})
			case opHeartbeatAck:
				sentHeartBeat = false
			case opReconnect:
				s.disconnect(websocket.CloseNormalClosure, "op Reconnect")
			case opDispatch:
				s.session.dispatchEvent(frame)
			default:
				logger.Errorf("Unknown opCode received: %d", frame.Op)
			}
		}
	}
}

func (s *shard) sendFrame(frame *gatewayFrame) {
	// Resume will only be sent within the reconnect lock, we don't relock to prevent a deadlock
	if frame.Op != opResume {
		s.lock.Lock()
		defer s.lock.Unlock()
	}

	logger.Debugf("Sending frame with opCode: %d", frame.Op)
	s.webSocket.WriteJSON(frame)
}

func (s *shard) readWebSocket(reader chan *receivedFrame) {
	logger.Debugf("Starting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	for {
		select {
		case <-s.stopRead:
			return
		default:
			frame, err := s.readFrame()
			if err != nil {
				if !s.session.isShuttingDown() {
					if closeErr, isClose := err.(*websocket.CloseError); isClose {
						s.closeMessage <- closeErr.Code
					}
					logger.ErrorE(err)
					go s.disconnect(websocket.CloseAbnormalClosure, err.Error())
				}
				<-s.stopRead // Wait for the connection to be closed
				return
			}

			reader <- frame
		}
	}

}

func (s *shard) readFrame() (*receivedFrame, error) {
	logger.Tracef("Shard.readFrame() called")
	msgType, msg, err := s.webSocket.ReadMessage()
	if err != nil {
		return nil, err
	}

	var reader io.Reader
	reader = bytes.NewBuffer(msg)

	if msgType == websocket.BinaryMessage {
		zReader, err := zlib.NewReader(reader)
		if err != nil {
			return nil, err
		}

		defer zReader.Close()

		reader = zReader
	}

	frame := receivedFrame{Sequence: -1}
	json.NewDecoder(reader).Decode(&frame)

	logger.Debugf("Received frame with opCode: %d", frame.Op)

	if frame.Sequence != -1 {
		logger.Tracef("Last sequence received set to %d.", frame.Sequence)
		s.sequence = frame.Sequence
	}
	return &frame, nil
}

func (s *shard) reconnect(recursive bool) {
	if !recursive {
		s.lock.Lock()
	}

	if !s.session.isShuttingDown() {
		logger.Noticef("Reconnecting shard [%d/%d]", s.shard+1, cap(s.session.shards))
		conn, _, err := websocket.DefaultDialer.Dial(s.session.wsUrl, http.Header{})

		if err == nil {
			s.webSocket = conn
			err = s.handshake()

			if err != nil {
				conn.Close()
				s.webSocket = nil
			}
		}

		if err != nil {
			logger.ErrorE(err)
			time.Sleep(1 * time.Second)
			go s.reconnect(true)
		} else {
			s.lock.Unlock()
		}
	}
}

func (s *shard) onClose(code int, text string) error {
	logger.Debugf("Received Close Frame from Discord. Code: %d. Text: %s", code, text)
	s.closeMessage <- code
	s.reconnect(false)
	return nil
}

func (s *shard) disconnect(code int, text string) {
	logger.Tracef("Shard.disconnect() called")
	s.stopListen <- true

	s.lock.Lock()
	err := s.webSocket.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, text))
	s.lock.Unlock()
	if err != nil {
		logger.ErrorE(err)
	}

	select {
	case code := <-s.closeMessage:
		logger.Debugf("Discord connection closed with code %d", code)
	case <-time.After(1 * time.Second):
		logger.Warn("Discord did not reply to the close message, force-closing connection.")
	}

	s.webSocket.Close()
	s.stopRead <- true

	go s.reconnect(false)
}
<|endoftext|>


<commit_before>package analogueclock

import (
	"math"
	"time"

	"github.com/kasworld/go-sdlgui"
	"github.com/kasworld/htmlcolors"
)

type Clock struct {
	*sdlgui.Control
	bg    htmlcolors.RGBA
	fg    htmlcolors.RGBA
	hhand htmlcolors.RGBA
	mhand htmlcolors.RGBA
	shand htmlcolors.RGBA
	t     time.Time
}

func New(x, y, z int, wx, wy int) *Clock {
	c := Clock{
		sdlgui.NewControl(x, y, z, wx, wy),
		htmlcolors.Black.ToRGBA(),
		htmlcolors.White.ToRGBA(),
		htmlcolors.Green.ToRGBA(),
		htmlcolors.Blue.ToRGBA(),
		htmlcolors.Red.ToRGBA(),
		time.Now(),
	}
	c.BorderSize = 2
	c.BorderType = 1
	return &c
}

func (tc *Clock) SetBG(bg htmlcolors.RGBA) {
	if bg == tc.bg {
		return
	}
	tc.ContentsChanged = true
	tc.bg = bg
}

func (tc *Clock) SetTime(t time.Time) {
	if t == tc.t {
		return
	}
	tc.ContentsChanged = true
	tc.t = t
}

func (tc *Clock) DrawFace() {
	tc.Rend.SetDrawColor(tc.fg[0], tc.fg[1], tc.fg[2], tc.fg[3])
	c := tc.GetRect().Center()
	rx := tc.GetRect().SizeVector()[0]/2 - 3
	ry := tc.GetRect().SizeVector()[1]/2 - 3
	rate := float64(ry) / float64(rx)
	for x := -rx; x < rx; x++ {
		theta := math.Acos(float64(x) / float64(rx))
		y := int(float64(x) * math.Tan(theta) * rate)
		if x == 0 {
			y = ry
		}
		tc.Rend.DrawLine(c[0]+x, c[1]-y, c[0]+x, c[1]+y)
	}
}

func (tc *Clock) DrawHand(co htmlcolors.RGBA, angle float64, l float64) {
	angle -= 90
	x1, y1 := tc.W/2, tc.H/2
	x2 := x1 + int(math.Cos(angle/180*math.Pi)*float64(x1)*l)
	y2 := y1 + int(math.Sin(angle/180*math.Pi)*float64(y1)*l)
	tc.Rend.SetDrawColor(co[0], co[1], co[2], co[3])
	tc.Rend.DrawLine(x1, y1, x2, y2)
}

func (tc *Clock) Time2Angle() (float64, float64, float64) {
	h := float64(tc.t.Hour()%12) / 12 * 360
	m := float64(tc.t.Minute()) / 60 * 360
	s := float64(tc.t.Second()) / 60 * 360
	return h, m, s
}

func (tc *Clock) DrawSurface() {
	if !tc.ContentsChanged {
		return
	}
	tc.ContentsChanged = false

	tc.Rend.SetDrawColor(tc.bg[0], tc.bg[1], tc.bg[2], tc.bg[3])
	tc.Rend.Clear()

	h, m, s := tc.Time2Angle()
	tc.DrawFace()
	tc.DrawHand(tc.hhand, h, 0.60)
	tc.DrawHand(tc.mhand, m, 0.75)
	tc.DrawHand(tc.shand, s, 0.90)

	tc.Rend.Present()
	tc.Win.AddUpdateControl(tc)
}
<commit_msg>rename to local funcions<commit_after>package analogueclock

import (
	"math"
	"time"

	"github.com/kasworld/go-sdlgui"
	"github.com/kasworld/htmlcolors"
)

type Clock struct {
	*sdlgui.Control
	bg    htmlcolors.RGBA
	fg    htmlcolors.RGBA
	hhand htmlcolors.RGBA
	mhand htmlcolors.RGBA
	shand htmlcolors.RGBA
	t     time.Time
}

func New(x, y, z int, wx, wy int) *Clock {
	c := Clock{
		sdlgui.NewControl(x, y, z, wx, wy),
		htmlcolors.Black.ToRGBA(),
		htmlcolors.White.ToRGBA(),
		htmlcolors.Green.ToRGBA(),
		htmlcolors.Blue.ToRGBA(),
		htmlcolors.Red.ToRGBA(),
		time.Now(),
	}
	c.BorderSize = 2
	c.BorderType = 1
	return &c
}

func (tc *Clock) SetBG(bg htmlcolors.RGBA) {
	if bg == tc.bg {
		return
	}
	tc.ContentsChanged = true
	tc.bg = bg
}

func (tc *Clock) SetTime(t time.Time) {
	if t == tc.t {
		return
	}
	tc.ContentsChanged = true
	tc.t = t
}

func (tc *Clock) drawFace() {
	tc.Rend.SetDrawColor(tc.fg[0], tc.fg[1], tc.fg[2], tc.fg[3])
	c := tc.GetRect().Center()
	rx := tc.GetRect().SizeVector()[0]/2 - 3
	ry := tc.GetRect().SizeVector()[1]/2 - 3
	rate := float64(ry) / float64(rx)
	for x := -rx; x < rx; x++ {
		theta := math.Acos(float64(x) / float64(rx))
		y := int(float64(x) * math.Tan(theta) * rate)
		if x == 0 {
			y = ry
		}
		tc.Rend.DrawLine(c[0]+x, c[1]-y, c[0]+x, c[1]+y)
	}
}

func (tc *Clock) drawHand(co htmlcolors.RGBA, angle float64, l float64) {
	angle -= 90
	x1, y1 := tc.W/2, tc.H/2
	x2 := x1 + int(math.Cos(angle/180*math.Pi)*float64(x1)*l)
	y2 := y1 + int(math.Sin(angle/180*math.Pi)*float64(y1)*l)
	tc.Rend.SetDrawColor(co[0], co[1], co[2], co[3])
	tc.Rend.DrawLine(x1, y1, x2, y2)
}

func (tc *Clock) time2Angle() (float64, float64, float64) {
	h := float64(tc.t.Hour()%12) / 12 * 360
	m := float64(tc.t.Minute()) / 60 * 360
	s := float64(tc.t.Second()) / 60 * 360
	return h, m, s
}

func (tc *Clock) DrawSurface() {
	if !tc.ContentsChanged {
		return
	}
	tc.ContentsChanged = false

	tc.Rend.SetDrawColor(tc.bg[0], tc.bg[1], tc.bg[2], tc.bg[3])
	tc.Rend.Clear()

	h, m, s := tc.time2Angle()
	tc.drawFace()
	tc.drawHand(tc.hhand, h, 0.60)
	tc.drawHand(tc.mhand, m, 0.75)
	tc.drawHand(tc.shand, s, 0.90)

	tc.Rend.Present()
	tc.Win.AddUpdateControl(tc)
}
<|endoftext|>