<commit_before>package scamp

import (
	"errors"
	"net"
	"crypto/tls"
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"bytes"
	"io/ioutil"
	"time"

	"sync"
	"sync/atomic"
)

// Two minute timeout on clients
var msgTimeout = time.Second * 120

type ServiceActionFunc func(*Message, *Client)
type ServiceAction struct {
	callback ServiceActionFunc
	crudTags string
	version  int
}

type Service struct {
	serviceSpec   string
	sector        string
	name          string
	humanName     string

	listener      net.Listener
	listenerIP    net.IP
	listenerPort  int

	actions       map[string]*ServiceAction
	isRunning     bool

	clientsM      sync.Mutex
	clients       []*Client

	// requests      ClientChan

	cert          tls.Certificate
	pemCert       []byte // just a copy of what was read off disk at tls cert load time

	// stats
	statsCloseChan chan bool
	connectionsAccepted uint64
}

func NewService(sector string, serviceSpec string, humanName string) (serv *Service, err error){
	if len(humanName) > 18 {
		err = fmt.Errorf("name `%s` is too long, must be less than 18 bytes", humanName)
		return
	}

	serv = new(Service)
	serv.sector = sector
	serv.serviceSpec = serviceSpec
	serv.humanName = humanName
	serv.generateRandomName()

	serv.actions = make(map[string]*ServiceAction)

	crtPath := defaultConfig.ServiceCertPath(serv.humanName)
	keyPath := defaultConfig.ServiceKeyPath(serv.humanName)

	if crtPath == nil || keyPath == nil {
		err = fmt.Errorf( "could not find valid crt/key pair for service %s (`%s`,`%s`)", serv.humanName, crtPath, keyPath )
		return
	}

	// Load keypair for tls socket library to use
	serv.cert, err = tls.LoadX509KeyPair( string(crtPath), string(keyPath) )
	if err != nil {
		return
	}

	// Load cert in to memory for announce packet writing
	serv.pemCert, err = ioutil.ReadFile(string(crtPath))
	if err != nil {
		return
	}
	serv.pemCert = bytes.TrimSpace(serv.pemCert)

	// Finally, get ready for incoming requests
	err = serv.listen()
	if err != nil {
		return
	}

	serv.statsCloseChan = make(chan bool)
	go PrintStatsLoop(serv, time.Duration(15)*time.Second, serv.statsCloseChan)

	Trace.Printf("done initializing service")

	return
}

// TODO: port discovery and interface/IP discovery should happen here
// important to set values so announce packets are correct
func (serv *Service)listen() (err error) {
	config := &tls.Config{
		Certificates: []tls.Certificate{ serv.cert },
	}

	Info.Printf("starting service on %s", serv.serviceSpec)
	serv.listener,err = tls.Listen("tcp", serv.serviceSpec, config)
	if err != nil {
		return err
	}
	addr := serv.listener.Addr()
	Info.Printf("service now listening to %s", addr.String())

  // TODO: get listenerIP to return 127.0.0.1 or something other than '::'/nil
  // serv.listenerIP = serv.listener.Addr().(*net.TCPAddr).IP
  serv.listenerIP, err = IPForAnnouncePacket()
  Trace.Printf("serv.listenerIP: `%s`", serv.listenerIP)
  Info.Printf("serv.listenerIP(before): %s", serv.listenerIP)
  serv.listenerIP = net.ParseIP("127.0.0.1"),
  // Info.Printf("serv.listenerIP(after): %s", serv.listenerIP)
  if err != nil {
  	return
  }

	serv.listenerPort = serv.listener.Addr().(*net.TCPAddr).Port

	return
}

func (serv *Service)Register(name string, callback ServiceActionFunc) (err error) {
	if serv.isRunning {
		err = errors.New("cannot register handlers while server is running")
		return
	}

	serv.actions[name] = &ServiceAction {
		callback: callback,
		version: 1,
	}

	return
}

func (serv *Service)Run() {

	forLoop:
	for {
		netConn,err := serv.listener.Accept()
		if err != nil {
			Info.Printf("exiting service service Run(): `%s`", err)
			break forLoop
		}
		Trace.Printf("accepted new connection...")

		var tlsConn (*tls.Conn) = (netConn).(*tls.Conn)
		if tlsConn == nil {
			Error.Fatalf("could not create tlsConn")
			break forLoop
		}

		conn := NewConnection(tlsConn,"service")
		client := NewClient(conn)

		serv.clientsM.Lock()
		serv.clients = append(serv.clients, client)
		serv.clientsM.Unlock()

		go serv.Handle(client)

		atomic.AddUint64(&serv.connectionsAccepted, 1)
	}

	Info.Printf("closing all registered objects")

	serv.clientsM.Lock()
	defer serv.clientsM.Unlock()
	for _,client := range serv.clients {
		client.Close()
	}

	serv.statsCloseChan <- true
}

func (serv *Service)Handle(client *Client) {
	var action *ServiceAction

	HandlerLoop:
	for {
		select {
		case msg,ok := <-client.Incoming():
			if !ok {
				break HandlerLoop
			}
			action = serv.actions[msg.Action]

			if action != nil{
				// yay
				action.callback(msg, client)
			} else {
				Error.Printf("do not know how to handle action `%s`", msg.Action)

				reply := NewMessage()
		    reply.SetMessageType(MESSAGE_TYPE_REPLY)
		    reply.SetEnvelope(ENVELOPE_JSON)
		    reply.SetRequestId(msg.RequestId)
		    reply.Write([]byte(`{"error": "no such action"`))
				_,err := client.Send(reply)
				if err != nil {
					client.Close()
					break HandlerLoop
				}

			}
		case <- time.After(msgTimeout):
			Error.Printf("timeout... dying!")
			break HandlerLoop
		}
	}

	client.Close()
	serv.RemoveClient(client)

	Trace.Printf("done handling client")

}

func (serv *Service)RemoveClient(client *Client) (err error){
	serv.clientsM.Lock()
	defer serv.clientsM.Unlock()

	index := -1
	for i,entry := range serv.clients {
		if client == entry {
			index = i
			break
		}
	}

	if index == -1 {
		Error.Printf("tried removing client that wasn't being tracked")
		return fmt.Errorf("unknown client") // TODO can I get the client's IP?
	}

	client.Close()
	serv.clients = append(serv.clients[:index], serv.clients[index+1:]...)

	return nil
}

func (serv *Service)Stop(){
	// Sometimes we Stop() before service after service has been init but before it is started
	// The usual case is a bad config in another plugin
	if serv.listener != nil {
		serv.listener.Close()
	}
}

func (serv *Service)MarshalText() (b []byte, err error){
	var buf bytes.Buffer

	serviceProxy := ServiceAsServiceProxy(serv)
	classRecord,err := json.Marshal(&serviceProxy)
	if err != nil {
		return
	}

	sig, err := SignSHA256(classRecord, serv.cert.PrivateKey.(*rsa.PrivateKey))
	if err != nil {
		return
	}
	sigParts := stringToRows(sig, 76)

	buf.Write(classRecord)
	buf.WriteString("\n\n")
	buf.Write(serv.pemCert)
	buf.WriteString("\n\n")
	// buf.WriteString(sig)
	// buf.WriteString("\n\n")
	for _,part := range sigParts {
		buf.WriteString(part)
		buf.WriteString("\n")
	}
	buf.WriteString("\n")

	b = buf.Bytes()
	return
}

func stringToRows(input string, rowlen int) (output []string) {
	output = make([]string,0)

	if len(input) <= 76 {
		output = append(output, input)
	} else {
		substr := input[:]
		var row string
		var done bool = false
		for {
			if len(substr) > 76 {
				row = substr[0:76]
				substr = substr[76:]
			} else {
				row = substr[:]
				done = true
			}
			output = append(output,row)
			if done {
				break
			}
		}
	}

	return
}

func (serv *Service)generateRandomName() {
	randBytes := make([]byte, 18, 18)
	read,err := rand.Read(randBytes)
	if err != nil {
		err = fmt.Errorf("could not generate all rand bytes needed. only read %d of 18", read)
		return
	}
	base64RandBytes := base64.StdEncoding.EncodeToString(randBytes)

	var buffer bytes.Buffer
	buffer.WriteString(serv.humanName)
	buffer.WriteString("-")
	buffer.WriteString(base64RandBytes[0:])
	serv.name = string(buffer.Bytes())
}
<commit_msg>debug<commit_after>package scamp

import (
	"errors"
	"net"
	"crypto/tls"
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"bytes"
	"io/ioutil"
	"time"

	"sync"
	"sync/atomic"
)

// Two minute timeout on clients
var msgTimeout = time.Second * 120

type ServiceActionFunc func(*Message, *Client)
type ServiceAction struct {
	callback ServiceActionFunc
	crudTags string
	version  int
}

type Service struct {
	serviceSpec   string
	sector        string
	name          string
	humanName     string

	listener      net.Listener
	listenerIP    net.IP
	listenerPort  int

	actions       map[string]*ServiceAction
	isRunning     bool

	clientsM      sync.Mutex
	clients       []*Client

	// requests      ClientChan

	cert          tls.Certificate
	pemCert       []byte // just a copy of what was read off disk at tls cert load time

	// stats
	statsCloseChan chan bool
	connectionsAccepted uint64
}

func NewService(sector string, serviceSpec string, humanName string) (serv *Service, err error){
	if len(humanName) > 18 {
		err = fmt.Errorf("name `%s` is too long, must be less than 18 bytes", humanName)
		return
	}

	serv = new(Service)
	serv.sector = sector
	serv.serviceSpec = serviceSpec
	serv.humanName = humanName
	serv.generateRandomName()

	serv.actions = make(map[string]*ServiceAction)

	crtPath := defaultConfig.ServiceCertPath(serv.humanName)
	keyPath := defaultConfig.ServiceKeyPath(serv.humanName)

	if crtPath == nil || keyPath == nil {
		err = fmt.Errorf( "could not find valid crt/key pair for service %s (`%s`,`%s`)", serv.humanName, crtPath, keyPath )
		return
	}

	// Load keypair for tls socket library to use
	serv.cert, err = tls.LoadX509KeyPair( string(crtPath), string(keyPath) )
	if err != nil {
		return
	}

	// Load cert in to memory for announce packet writing
	serv.pemCert, err = ioutil.ReadFile(string(crtPath))
	if err != nil {
		return
	}
	serv.pemCert = bytes.TrimSpace(serv.pemCert)

	// Finally, get ready for incoming requests
	err = serv.listen()
	if err != nil {
		return
	}

	serv.statsCloseChan = make(chan bool)
	go PrintStatsLoop(serv, time.Duration(15)*time.Second, serv.statsCloseChan)

	Trace.Printf("done initializing service")

	return
}

// TODO: port discovery and interface/IP discovery should happen here
// important to set values so announce packets are correct
func (serv *Service)listen() (err error) {
	config := &tls.Config{
		Certificates: []tls.Certificate{ serv.cert },
	}

	Info.Printf("starting service on %s", serv.serviceSpec)
	serv.listener,err = tls.Listen("tcp", serv.serviceSpec, config)
	if err != nil {
		return err
	}
	addr := serv.listener.Addr()
	Info.Printf("service now listening to %s", addr.String())

  // TODO: get listenerIP to return 127.0.0.1 or something other than '::'/nil
  // serv.listenerIP = serv.listener.Addr().(*net.TCPAddr).IP
  serv.listenerIP, err = IPForAnnouncePacket()
  Trace.Printf("serv.listenerIP: `%s`", serv.listenerIP)
  Info.Printf("serv.listenerIP(before): %s", serv.listenerIP)
  serv.listenerIP = net.ParseIP("127.0.0.1")
  Info.Printf("serv.listenerIP(after): %s", serv.listenerIP)
  if err != nil {
  	return
  }

	serv.listenerPort = serv.listener.Addr().(*net.TCPAddr).Port

	return
}

func (serv *Service)Register(name string, callback ServiceActionFunc) (err error) {
	if serv.isRunning {
		err = errors.New("cannot register handlers while server is running")
		return
	}

	serv.actions[name] = &ServiceAction {
		callback: callback,
		version: 1,
	}

	return
}

func (serv *Service)Run() {

	forLoop:
	for {
		netConn,err := serv.listener.Accept()
		if err != nil {
			Info.Printf("exiting service service Run(): `%s`", err)
			break forLoop
		}
		Trace.Printf("accepted new connection...")

		var tlsConn (*tls.Conn) = (netConn).(*tls.Conn)
		if tlsConn == nil {
			Error.Fatalf("could not create tlsConn")
			break forLoop
		}

		conn := NewConnection(tlsConn,"service")
		client := NewClient(conn)

		serv.clientsM.Lock()
		serv.clients = append(serv.clients, client)
		serv.clientsM.Unlock()

		go serv.Handle(client)

		atomic.AddUint64(&serv.connectionsAccepted, 1)
	}

	Info.Printf("closing all registered objects")

	serv.clientsM.Lock()
	defer serv.clientsM.Unlock()
	for _,client := range serv.clients {
		client.Close()
	}

	serv.statsCloseChan <- true
}

func (serv *Service)Handle(client *Client) {
	var action *ServiceAction

	HandlerLoop:
	for {
		select {
		case msg,ok := <-client.Incoming():
			if !ok {
				break HandlerLoop
			}
			action = serv.actions[msg.Action]

			if action != nil{
				// yay
				action.callback(msg, client)
			} else {
				Error.Printf("do not know how to handle action `%s`", msg.Action)

				reply := NewMessage()
		    reply.SetMessageType(MESSAGE_TYPE_REPLY)
		    reply.SetEnvelope(ENVELOPE_JSON)
		    reply.SetRequestId(msg.RequestId)
		    reply.Write([]byte(`{"error": "no such action"`))
				_,err := client.Send(reply)
				if err != nil {
					client.Close()
					break HandlerLoop
				}

			}
		case <- time.After(msgTimeout):
			Error.Printf("timeout... dying!")
			break HandlerLoop
		}
	}

	client.Close()
	serv.RemoveClient(client)

	Trace.Printf("done handling client")

}

func (serv *Service)RemoveClient(client *Client) (err error){
	serv.clientsM.Lock()
	defer serv.clientsM.Unlock()

	index := -1
	for i,entry := range serv.clients {
		if client == entry {
			index = i
			break
		}
	}

	if index == -1 {
		Error.Printf("tried removing client that wasn't being tracked")
		return fmt.Errorf("unknown client") // TODO can I get the client's IP?
	}

	client.Close()
	serv.clients = append(serv.clients[:index], serv.clients[index+1:]...)

	return nil
}

func (serv *Service)Stop(){
	// Sometimes we Stop() before service after service has been init but before it is started
	// The usual case is a bad config in another plugin
	if serv.listener != nil {
		serv.listener.Close()
	}
}

func (serv *Service)MarshalText() (b []byte, err error){
	var buf bytes.Buffer

	serviceProxy := ServiceAsServiceProxy(serv)
	classRecord,err := json.Marshal(&serviceProxy)
	if err != nil {
		return
	}

	sig, err := SignSHA256(classRecord, serv.cert.PrivateKey.(*rsa.PrivateKey))
	if err != nil {
		return
	}
	sigParts := stringToRows(sig, 76)

	buf.Write(classRecord)
	buf.WriteString("\n\n")
	buf.Write(serv.pemCert)
	buf.WriteString("\n\n")
	// buf.WriteString(sig)
	// buf.WriteString("\n\n")
	for _,part := range sigParts {
		buf.WriteString(part)
		buf.WriteString("\n")
	}
	buf.WriteString("\n")

	b = buf.Bytes()
	return
}

func stringToRows(input string, rowlen int) (output []string) {
	output = make([]string,0)

	if len(input) <= 76 {
		output = append(output, input)
	} else {
		substr := input[:]
		var row string
		var done bool = false
		for {
			if len(substr) > 76 {
				row = substr[0:76]
				substr = substr[76:]
			} else {
				row = substr[:]
				done = true
			}
			output = append(output,row)
			if done {
				break
			}
		}
	}

	return
}

func (serv *Service)generateRandomName() {
	randBytes := make([]byte, 18, 18)
	read,err := rand.Read(randBytes)
	if err != nil {
		err = fmt.Errorf("could not generate all rand bytes needed. only read %d of 18", read)
		return
	}
	base64RandBytes := base64.StdEncoding.EncodeToString(randBytes)

	var buffer bytes.Buffer
	buffer.WriteString(serv.humanName)
	buffer.WriteString("-")
	buffer.WriteString(base64RandBytes[0:])
	serv.name = string(buffer.Bytes())
}
<|endoftext|>


<commit_before>package mrT

import (
	"bytes"
	"fmt"
	"os"
	"testing"

	"github.com/missionMeteora/journaler"
)

const (
	testInvalidActionsFmt = "invalid number of actions, expected %d and received %d"
)

func TestMrT(t *testing.T) {
	var (
		m        *MrT
		firstTxn string
		err      error
	)

	if m, err = New("./testing/", "testing"); err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll("./testing/")

	if err = m.Txn(func(txn *Txn) (err error) {
		txn.Put([]byte("greeting"), []byte("hello"))
		txn.Put([]byte("name"), []byte("world"))
		return
	}); err != nil {
		return
	}

	if err = m.Txn(func(txn *Txn) (err error) {
		txn.Put([]byte("name"), []byte("John Doe"))
		return
	}); err != nil {
		return
	}

	if firstTxn, err = peekFirstTxn(m.s); err != nil {
		t.Fatal(err)
	}

	if err = testForEach(m, "", 3); err != nil {
		t.Fatal(err)
	}

	if err = testForEachTxn(m, "", 2); err != nil {
		t.Fatal(err)
	}

	if err = testForEach(m, firstTxn, 1); err != nil {
		t.Fatal(err)
	}

	if err = m.Archive(func(txn *Txn) (err error) {
		txn.Put([]byte("greeting"), []byte("hello"))
		txn.Put([]byte("name"), []byte("John Doe"))
		return
	}); err != nil {
		t.Fatal(err)
	}

	if err = testForEach(m, firstTxn, 1); err != nil {
		t.Fatal(err)
	}

	if err = testForEachTxn(m, "", 2); err != nil {
		t.Fatal(err)
	}

	if err = m.Txn(func(txn *Txn) (err error) {
		txn.Put([]byte("name"), []byte("derp"))
		return
	}); err != nil {
		return
	}

	if err = testForEach(m, "", 4); err != nil {
		t.Fatal(err)
	}

	if err = testForEachTxn(m, "", 3); err != nil {
		t.Fatal(err)
	}

	if err = m.Archive(func(txn *Txn) (err error) {
		txn.Put([]byte("greeting"), []byte("hello"))
		txn.Put([]byte("name"), []byte("derp"))
		return
	}); err != nil {
		t.Fatal(err)
	}

	if err = testForEach(m, "", 4); err != nil {
		t.Fatal(err)
	}

	if err = testForEachTxn(m, "", 3); err != nil {
		t.Fatal(err)
	}

	buf := bytes.NewBuffer(nil)
	if err = m.Export("", buf); err != nil {
		t.Fatal(err)
	}

	var nm *MrT
	if nm, err = New("./testing2/", "testing"); err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll("./testing2/")

	var lastTxn string
	if lastTxn, err = nm.Import(buf, testNilForEach); err != nil {
		t.Fatal(err)
	}

	// We will only expect two entries because we will pull the "current" data
	if err = testForEach(nm, "", 2); err != nil {
		t.Fatal(err)
	}

	// We will only expect one transaction because we will pull the "current" data
	if err = testForEachTxn(nm, "", 1); err != nil {
		t.Fatal(err)
	}

	if err = m.Txn(func(txn *Txn) (err error) {
		txn.Put([]byte("name"), []byte("foo"))
		return
	}); err != nil {
		t.Fatal(err)
	}

	if err = m.Export(lastTxn, buf); err != nil {
		t.Fatal(err)
	}

	if lastTxn, err = nm.Import(buf, testNilForEach); err != nil {
		t.Fatal(err)
	}

	// We will only expect two entries because we will pull the "current" data
	if err = testForEach(nm, "", 3); err != nil {
		t.Fatal(err)
	}

	journaler.Success("Testing complete")
	return
}

func testNilForEach(lineType byte, key, value []byte) (err error) {
	return
}

func testForEach(m *MrT, start string, n int) (err error) {
	var entryCount int
	if err = m.ForEach(start, true, func(lineType byte, key []byte, value []byte) (err error) {
		if lineType != PutLine && lineType != DeleteLine {
			return
		}

		entryCount++
		return
	}); err != nil {
		return
	}

	if entryCount != n {
		return fmt.Errorf("invalid number of entries, expected %d and recieved %d", n, entryCount)
	}

	return
}

func testForEachTxn(m *MrT, start string, n int) (err error) {
	var entryCount int
	if err = m.ForEachTxn(start, true, func(ti *TxnInfo) (err error) {
		switch entryCount {
		case 0:
			if len(ti.Actions) != 2 {
				return fmt.Errorf(testInvalidActionsFmt, 2, len(ti.Actions))
			}

		case 1:
			if len(ti.Actions) != 1 {
				return fmt.Errorf(testInvalidActionsFmt, 1, len(ti.Actions))
			}

		case 2:
			if len(ti.Actions) != 1 {
				return fmt.Errorf(testInvalidActionsFmt, 1, len(ti.Actions))
			}

		default:
			return fmt.Errorf("invalid number of entries, recieved %d", entryCount)
		}

		entryCount++
		return
	}); err != nil {
		return
	}

	if entryCount != n {
		return fmt.Errorf("invalid number of entries, expected %d and recieved %d", n, entryCount)
	}

	return
}
<commit_msg>Adjust tests to accomodate for new approach to handling empty txn id's<commit_after>package mrT

import (
	"bytes"
	"fmt"
	"os"
	"testing"

	"github.com/missionMeteora/journaler"
)

const (
	testInvalidActionsFmt = "invalid number of actions, expected %d and received %d"
)

func TestMrT(t *testing.T) {
	var (
		m        *MrT
		firstTxn string
		err      error
	)

	if m, err = New("./testing/", "testing"); err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll("./testing/")

	if err = m.Txn(func(txn *Txn) (err error) {
		txn.Put([]byte("greeting"), []byte("hello"))
		txn.Put([]byte("name"), []byte("world"))
		return
	}); err != nil {
		return
	}

	if err = m.Txn(func(txn *Txn) (err error) {
		txn.Put([]byte("name"), []byte("John Doe"))
		return
	}); err != nil {
		return
	}

	if firstTxn, err = peekFirstTxn(m.s); err != nil {
		t.Fatal(err)
	}

	if err = testForEach(m, "", 3); err != nil {
		t.Fatal(err)
	}

	if err = testForEachTxn(m, "", 2); err != nil {
		t.Fatal(err)
	}

	if err = testForEach(m, firstTxn, 1); err != nil {
		t.Fatal(err)
	}

	if err = m.Archive(func(txn *Txn) (err error) {
		txn.Put([]byte("greeting"), []byte("hello"))
		txn.Put([]byte("name"), []byte("John Doe"))
		return
	}); err != nil {
		t.Fatal(err)
	}

	if err = testForEach(m, firstTxn, 1); err != nil {
		t.Fatal(err)
	}

	if err = testForEachTxn(m, "", 1); err != nil {
		t.Fatal(err)
	}

	if err = m.Txn(func(txn *Txn) (err error) {
		txn.Put([]byte("name"), []byte("derp"))
		return
	}); err != nil {
		return
	}

	if err = testForEach(m, "", 3); err != nil {
		t.Fatal(err)
	}

	if err = testForEachTxn(m, "", 2); err != nil {
		t.Fatal(err)
	}

	if err = m.Archive(func(txn *Txn) (err error) {
		txn.Put([]byte("greeting"), []byte("hello"))
		txn.Put([]byte("name"), []byte("derp"))
		return
	}); err != nil {
		t.Fatal(err)
	}

	if err = testForEach(m, "", 2); err != nil {
		t.Fatal(err)
	}

	if err = testForEachTxn(m, "", 1); err != nil {
		t.Fatal(err)
	}

	buf := bytes.NewBuffer(nil)
	if err = m.Export("", buf); err != nil {
		t.Fatal(err)
	}

	var nm *MrT
	if nm, err = New("./testing2/", "testing"); err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll("./testing2/")

	var lastTxn string
	if lastTxn, err = nm.Import(buf, testNilForEach); err != nil {
		t.Fatal(err)
	}

	// We will only expect two entries because we will pull the "current" data
	if err = testForEach(nm, "", 2); err != nil {
		t.Fatal(err)
	}

	// We will only expect one transaction because we will pull the "current" data
	if err = testForEachTxn(nm, "", 1); err != nil {
		t.Fatal(err)
	}

	if err = m.Txn(func(txn *Txn) (err error) {
		txn.Put([]byte("name"), []byte("foo"))
		return
	}); err != nil {
		t.Fatal(err)
	}

	if err = m.Export(lastTxn, buf); err != nil {
		t.Fatal(err)
	}

	if lastTxn, err = nm.Import(buf, testNilForEach); err != nil {
		t.Fatal(err)
	}

	// We will only expect two entries because we will pull the "current" data
	if err = testForEach(nm, "", 3); err != nil {
		t.Fatal(err)
	}

	journaler.Success("Testing complete")
	return
}

func testNilForEach(lineType byte, key, value []byte) (err error) {
	return
}

func testForEach(m *MrT, start string, n int) (err error) {
	var entryCount int
	if err = m.ForEach(start, true, func(lineType byte, key []byte, value []byte) (err error) {
		if lineType != PutLine && lineType != DeleteLine {
			return
		}

		entryCount++
		return
	}); err != nil {
		return
	}

	if entryCount != n {
		return fmt.Errorf("invalid number of entries, expected %d and recieved %d", n, entryCount)
	}

	return
}

func testForEachTxn(m *MrT, start string, n int) (err error) {
	var entryCount int
	if err = m.ForEachTxn(start, true, func(ti *TxnInfo) (err error) {
		switch entryCount {
		case 0:
			if len(ti.Actions) != 2 {
				return fmt.Errorf(testInvalidActionsFmt, 2, len(ti.Actions))
			}

		case 1:
			if len(ti.Actions) != 1 {
				return fmt.Errorf(testInvalidActionsFmt, 1, len(ti.Actions))
			}

		case 2:
			if len(ti.Actions) != 1 {
				return fmt.Errorf(testInvalidActionsFmt, 1, len(ti.Actions))
			}

		default:
			return fmt.Errorf("invalid number of entries, recieved %d", entryCount)
		}

		entryCount++
		return
	}); err != nil {
		return
	}

	if entryCount != n {
		return fmt.Errorf("invalid number of entries, expected %d and recieved %d", n, entryCount)
	}

	return
}
<|endoftext|>


<commit_before>package main

import (

  "log"
  "time"
  "flag"
  "github.com/bwmarrin/dgvoice"
  "github.com/bwmarrin/discordgo"
  "github.com/spf13/viper"
  "github.com/fsnotify/fsnotify"
  //"reflect"
  //"github.com/BurntSushi/toml"
)

const (
  channels  int = 2                   // 1 for mono, 2 for stereo
  frameRate int = 48000               // audio sampling rate
  frameSize int = 960                 // uint16 size of each audio frame
  //maxBytes  int = (frameSize * 2) * 2 // max size of opus data
)

var bot map[string]string

// loadConfig load the config file
func loadConfig (filename string) (error){
  // Read the config.toml file
  viper.SetConfigType("toml")
  viper.SetConfigFile(filename)
  log.Println("Opening", filename)
  //viper.AddConfigPath(".")
  err := viper.ReadInConfig()
  if err != nil {
    log.Println("Config file (bot.toml) not found.\nExit ...\n")
    return err
  }
  bot = viper.GetStringMapString("bot")
  log.Println("URL:", bot["url"])
  return nil
}

// connectionOn open the url stream connection
func connectionOn(filename string) {
  // load the config file
  err := loadConfig(filename)
  if err != nil {
    log.Println(err)
    return
  }

  discord, err := discordgo.New("Bot " + bot["token"])
  if err != nil {
    log.Println("error creating Discord session,", err)
    return
  }
  defer discord.Close()

  log.Println("Bot is Opening.")

  //discord.AddHandler(messageCreate)

  // Open Websocket
  err = discord.Open()
  if err != nil {
    log.Println("Error Open():", err)
    return
  }

  log.Println("Bot user test.")

  _, err = discord.User("@me")
  if err != nil {
    // Login unsuccessful
    log.Println(err)
    return
  } // Login successful

  log.Println("Bot is now running. Press CTRL-C to exit.")

  // Set Status
  discord.UpdateStatus(0, bot["status"])

  // Join to voice channel
  vs, err := discord.ChannelVoiceJoin(bot["guild_id"], bot["channel_id"], false, false)
  if err != nil {
    log.Println(err)
    return
  }
  defer vs.Close()

  // Play to URL
  dgvoice.PlayAudioFile(vs, bot["url"])
}

func main() {
  file_name := flag.String("f", "bot.toml", "Set path for the config file.")
  flag.Parse()

  // Hot reload
  viper.WatchConfig()
  viper.OnConfigChange(func (e fsnotify.Event) {
    log.Println("The config file changed:", e.Name)
    dgvoice.KillPlayer()
  })

  for {
    connectionOn(*file_name)
    // Restart the bot if the file is changed or the connection fail
    log.Println("The bot is restarting ...")
    time.Sleep(5000 * time.Millisecond)
  }

  return
}

<commit_msg>Remove code without use<commit_after>package main

import (

  "log"
  "time"
  "flag"
  "github.com/bwmarrin/dgvoice"
  "github.com/bwmarrin/discordgo"
  "github.com/spf13/viper"
  "github.com/fsnotify/fsnotify"
)

var bot map[string]string

// loadConfig load the config file
func loadConfig (filename string) (error){
  // Read the config.toml file
  viper.SetConfigType("toml")
  viper.SetConfigFile(filename)
  log.Println("Opening", filename)
  //viper.AddConfigPath(".")
  err := viper.ReadInConfig()
  if err != nil {
    log.Println("Config file (bot.toml) not found.\nExit ...\n")
    return err
  }
  bot = viper.GetStringMapString("bot")
  log.Println("URL:", bot["url"])
  return nil
}

// connectionOn open the url stream connection
func connectionOn(filename string) {
  // load the config file
  err := loadConfig(filename)
  if err != nil {
    log.Println(err)
    return
  }

  discord, err := discordgo.New("Bot " + bot["token"])
  if err != nil {
    log.Println("error creating Discord session,", err)
    return
  }
  defer discord.Close()

  log.Println("Bot is Opening.")

  //discord.AddHandler(messageCreate)

  // Open Websocket
  err = discord.Open()
  if err != nil {
    log.Println("Error Open():", err)
    return
  }

  log.Println("Bot user test.")

  _, err = discord.User("@me")
  if err != nil {
    // Login unsuccessful
    log.Println(err)
    return
  } // Login successful

  log.Println("Bot is now running. Press CTRL-C to exit.")

  // Set Status
  discord.UpdateStatus(0, bot["status"])

  // Join to voice channel
  vs, err := discord.ChannelVoiceJoin(bot["guild_id"], bot["channel_id"], false, false)
  if err != nil {
    log.Println(err)
    return
  }
  defer vs.Close()

  // Play to URL
  dgvoice.PlayAudioFile(vs, bot["url"])
}

func main() {
  file_name := flag.String("f", "bot.toml", "Set path for the config file.")
  flag.Parse()

  // Hot reload
  viper.WatchConfig()
  viper.OnConfigChange(func (e fsnotify.Event) {
    log.Println("The config file changed:", e.Name)
    dgvoice.KillPlayer()
  })

  for {
    connectionOn(*file_name)
    // Restart the bot if the file is changed or the connection fail
    log.Println("The bot is restarting ...")
    time.Sleep(5000 * time.Millisecond)
  }

  return
}

<|endoftext|>


<commit_before>package net

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"
)

func HTTPGet(url string, headers map[string]string, cookies []*http.Cookie) ([]byte, error) {
	return HTTPGetWithTimeout(url, headers, cookies, time.Duration(0))
}

func HTTPGetWithTimeout(url string, headers map[string]string, cookies []*http.Cookie, timeout time.Duration) ([]byte, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	for name, value := range headers {
		req.Header.Set(name, value)
	}

	for _, c := range cookies {
		req.AddCookie(c)
	}
	client := http.Client{
		Timeout: timeout,
	}
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return nil, fmt.Errorf("Failed to call [%s], status code: %d", url, res.StatusCode)
	}

	return ioutil.ReadAll(res.Body)
}

func HTTPPost(url string, bodyData []byte) ([]byte, error) {
	b := bytes.NewReader(bodyData)
	res, err := http.Post(url, "plain/text", b)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 && res.StatusCode != 202 {
		return nil, fmt.Errorf("Failed to call [%s], status code: %d", url, res.StatusCode)
	}

	return ioutil.ReadAll(res.Body)
}

func HTTPPut(url string, bodyData []byte) ([]byte, error) {
	b := bytes.NewReader(bodyData)
	client := &http.Client{}
	req, err := http.NewRequest("PUT", url, b)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "plain/text")
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 && resp.StatusCode != 202 {
		return nil, fmt.Errorf("Failed to call [%s], status code: %d", url, resp.StatusCode)
	}

	return ioutil.ReadAll(resp.Body)
}

func HTTPParseURL(srcURL string, params map[string][]string) (string, error) {
	u, err := url.Parse(srcURL)
	if err != nil {
		return "", fmt.Errorf("Failed to parse url: %s", err.Error())
	}
	p := url.Values{}
	for k, v := range params {
		for _, param := range v {
			if len(p.Get(k)) > 0 {
				p.Add(k, param)
			} else {
				p.Set(k, param)
			}
		}
	}
	u.RawQuery = p.Encode()
	return u.String(), nil
}

// FormPost do a HTTP POST with x-www-form-urlencoded body
func FormPost(url string, bodyData []byte) ([]byte, error) {
	b := bytes.NewReader(bodyData)
	req, err := http.NewRequest("POST", url, b)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	client := &http.Client{}
	res, err := client.Do(req)
	defer res.Body.Close()
	if err != nil {
		return nil, err
	}

	if res.StatusCode != 200 && res.StatusCode != 202 {
		return nil, fmt.Errorf("Failed to call [%s], status code: %d", url, res.StatusCode)
	}

	return ioutil.ReadAll(res.Body)
}
<commit_msg>Add HTTPDelete<commit_after>package net

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"
)

func HTTPGet(url string, headers map[string]string, cookies []*http.Cookie) ([]byte, error) {
	return HTTPGetWithTimeout(url, headers, cookies, time.Duration(0))
}

func CustomRequest(method string, url string, bodyData []byte) ([]byte, error) {
	b := bytes.NewReader(bodyData)
	client := &http.Client{}
	req, err := http.NewRequest(method, url, b)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "plain/text")
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 && resp.StatusCode != 202 {
		return nil, fmt.Errorf("Failed to call [%s], status code: %d", url, resp.StatusCode)
	}

	return ioutil.ReadAll(resp.Body)
}

func HTTPGetWithTimeout(url string, headers map[string]string, cookies []*http.Cookie, timeout time.Duration) ([]byte, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	for name, value := range headers {
		req.Header.Set(name, value)
	}

	for _, c := range cookies {
		req.AddCookie(c)
	}
	client := http.Client{
		Timeout: timeout,
	}
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return nil, fmt.Errorf("Failed to call [%s], status code: %d", url, res.StatusCode)
	}

	return ioutil.ReadAll(res.Body)
}

func HTTPPost(url string, bodyData []byte) ([]byte, error) {
	b := bytes.NewReader(bodyData)
	res, err := http.Post(url, "plain/text", b)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 && res.StatusCode != 202 {
		return nil, fmt.Errorf("Failed to call [%s], status code: %d", url, res.StatusCode)
	}

	return ioutil.ReadAll(res.Body)
}

func HTTPDelete(url string, bodyData []byte) ([]byte, error) {
	return CustomRequest("DELETE", url, bodyData)
}

func HTTPPut(url string, bodyData []byte) ([]byte, error) {
	return CustomRequest("PUT", url, bodyData)
}

func HTTPParseURL(srcURL string, params map[string][]string) (string, error) {
	u, err := url.Parse(srcURL)
	if err != nil {
		return "", fmt.Errorf("Failed to parse url: %s", err.Error())
	}
	p := url.Values{}
	for k, v := range params {
		for _, param := range v {
			if len(p.Get(k)) > 0 {
				p.Add(k, param)
			} else {
				p.Set(k, param)
			}
		}
	}
	u.RawQuery = p.Encode()
	return u.String(), nil
}

// FormPost do a HTTP POST with x-www-form-urlencoded body
func FormPost(url string, bodyData []byte) ([]byte, error) {
	b := bytes.NewReader(bodyData)
	req, err := http.NewRequest("POST", url, b)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	client := &http.Client{}
	res, err := client.Do(req)
	defer res.Body.Close()
	if err != nil {
		return nil, err
	}

	if res.StatusCode != 200 && res.StatusCode != 202 {
		return nil, fmt.Errorf("Failed to call [%s], status code: %d", url, res.StatusCode)
	}

	return ioutil.ReadAll(res.Body)
}
<|endoftext|>


<commit_before>package npc

import (
	"fmt"
	"strings"

	"k8s.io/apimachinery/pkg/types"

	"github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/npc/iptables"
)

type ruleHost interface {
	// getRuleSpec returns source or destination specification and comment which
	// are used in an iptables rule.
	//
	// If src=true then the rulespec is for source, otherwise - for destination.
	getRuleSpec(src bool) ([]string, string)
	// TODO(brb) Get rid
	isNil() bool
}

type ruleSpec struct {
	key        string
	args       []string
	policyType policyType
}

func newRuleSpec(policyType policyType, proto *string, srcHost, dstHost ruleHost, dstPort *string) *ruleSpec {
	args := []string{}
	if proto != nil {
		args = append(args, "-p", *proto)
	}
	srcComment := "anywhere"
	if srcHost != nil && !srcHost.isNil() {
		rule, comment := srcHost.getRuleSpec(true)
		args = append(args, rule...)
		srcComment = comment
	}
	dstComment := "anywhere"
	if dstHost != nil && !dstHost.isNil() {
		rule, comment := dstHost.getRuleSpec(false)
		args = append(args, rule...)
		dstComment = comment
	}
	if dstPort != nil {
		args = append(args, "--dport", *dstPort)
	}
	args = append(args, "-j", "ACCEPT")
	// NOTE: if you remove the comment bellow, then embed `policyType` into `key`.
	// Otherwise, the rule won't be provisioned if it exists for other policy type.
	args = append(args, "-m", "comment", "--comment", fmt.Sprintf("%s -> %s (%s)", srcComment, dstComment, policyTypeStr(policyType)))
	key := strings.Join(args, " ")

	return &ruleSpec{key, args, policyType}
}

type ruleSet struct {
	ipt   iptables.Interface
	users map[string]map[types.UID]struct{}
}

func newRuleSet(ipt iptables.Interface) *ruleSet {
	return &ruleSet{ipt, make(map[string]map[types.UID]struct{})}
}

func (rs *ruleSet) deprovision(user types.UID, current, desired map[string]*ruleSpec) error {
	for key, spec := range current {
		if _, found := desired[key]; !found {
			delete(rs.users[key], user)
			if len(rs.users[key]) == 0 {
				chain := rulesChainByPolicyType(spec.policyType)
				common.Log.Infof("deleting rule %v from %q chain", spec.args, chain)
				if err := rs.ipt.Delete(TableFilter, chain, spec.args...); err != nil {
					return err
				}
				delete(rs.users, key)
			}
		}
	}

	return nil
}

func (rs *ruleSet) provision(user types.UID, current, desired map[string]*ruleSpec) error {
	for key, spec := range desired {
		if _, found := current[key]; !found {
			if _, found := rs.users[key]; !found {
				chain := rulesChainByPolicyType(spec.policyType)
				common.Log.Infof("adding rule %v to %q chain", spec.args, chain)
				if err := rs.ipt.Append(TableFilter, chain, spec.args...); err != nil {
					return err
				}
				rs.users[key] = make(map[types.UID]struct{})
			}
			rs.users[key][user] = struct{}{}
		}
	}

	return nil
}

func rulesChainByPolicyType(policyType policyType) string {
	if policyType == egressPolicy {
		return EgressCustomChain
	}
	return IngressChain
}
<commit_msg>npc: Provision egress rules<commit_after>package npc

import (
	"fmt"
	"strings"

	"k8s.io/apimachinery/pkg/types"

	"github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/npc/iptables"
)

type ruleHost interface {
	// getRuleSpec returns source or destination specification and comment which
	// are used in an iptables rule.
	//
	// If src=true then the rulespec is for source, otherwise - for destination.
	getRuleSpec(src bool) ([]string, string)
	// TODO(brb) Get rid
	isNil() bool
}

type ruleSpec struct {
	key        string
	args       []string
	policyType policyType
}

func newRuleSpec(policyType policyType, proto *string, srcHost, dstHost ruleHost, dstPort *string) *ruleSpec {
	args := []string{}
	if proto != nil {
		args = append(args, "-p", *proto)
	}
	srcComment := "anywhere"
	if srcHost != nil && !srcHost.isNil() {
		rule, comment := srcHost.getRuleSpec(true)
		args = append(args, rule...)
		srcComment = comment
	}
	dstComment := "anywhere"
	if dstHost != nil && !dstHost.isNil() {
		rule, comment := dstHost.getRuleSpec(false)
		args = append(args, rule...)
		dstComment = comment
	}
	if dstPort != nil {
		args = append(args, "--dport", *dstPort)
	}
	// NOTE: if you remove the comment bellow, then embed `policyType` into `key`.
	// Otherwise, the rule won't be provisioned if it exists for other policy type.
	args = append(args, "-m", "comment", "--comment", fmt.Sprintf("%s -> %s (%s)", srcComment, dstComment, policyTypeStr(policyType)))
	key := strings.Join(args, " ")

	return &ruleSpec{key, args, policyType}
}

func (spec *ruleSpec) iptChain() string {
	if spec.policyType == egressPolicy {
		return EgressCustomChain
	}
	return IngressChain
}

func (spec *ruleSpec) iptRuleSpecs() [][]string {
	if spec.policyType == ingressPolicy {
		rule := make([]string, len(spec.args))
		copy(rule, spec.args)
		rule = append(rule, "-j", "ACCEPT")
		return [][]string{rule}
	}

	// egressPolicy
	ruleMark := make([]string, len(spec.args))
	copy(ruleMark, spec.args)
	ruleMark = append(ruleMark, "-j", "MARK", "--set-xmark", EgressMark)
	ruleReturn := make([]string, len(spec.args))
	copy(ruleReturn, spec.args)
	ruleReturn = append(ruleReturn, "-j", "RETURN")
	return [][]string{ruleMark, ruleReturn}
}

type ruleSet struct {
	ipt   iptables.Interface
	users map[string]map[types.UID]struct{}
}

func newRuleSet(ipt iptables.Interface) *ruleSet {
	return &ruleSet{ipt, make(map[string]map[types.UID]struct{})}
}

func (rs *ruleSet) deprovision(user types.UID, current, desired map[string]*ruleSpec) error {
	for key, spec := range current {
		if _, found := desired[key]; !found {
			delete(rs.users[key], user)
			if len(rs.users[key]) == 0 {
				chain := spec.iptChain()
				for _, rule := range spec.iptRuleSpecs() {
					common.Log.Infof("deleting rule %v from %q chain", rule, chain)
					if err := rs.ipt.Delete(TableFilter, chain, rule...); err != nil {
						return err
					}
				}

				delete(rs.users, key)
			}
		}
	}

	return nil
}

func (rs *ruleSet) provision(user types.UID, current, desired map[string]*ruleSpec) error {
	for key, spec := range desired {
		if _, found := current[key]; !found {
			if _, found := rs.users[key]; !found {
				chain := spec.iptChain()
				for _, rule := range spec.iptRuleSpecs() {
					common.Log.Infof("adding rule %v to %q chain", rule, chain)
					if err := rs.ipt.Append(TableFilter, chain, rule...); err != nil {
						return err
					}
				}
				rs.users[key] = make(map[types.UID]struct{})
			}
			rs.users[key][user] = struct{}{}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package orm

import (
	"context"
	"database/sql"
	"fmt"
	"os"

	"github.com/pkg/errors"
	"github.com/rbastic/dyndao/object"
	"github.com/rbastic/dyndao/schema"
)

// GetParentsViaChild retrieves all direct (one-level 'up') parents for a given child object.
// If a child contains multiple parent tables (possibility?) then this would return an Array
// of objects with multiple potential values for their obj.Type fields.
func (o ORM) GetParentsViaChild(ctx context.Context, childObj *object.Object) (object.Array, error) {
	table := childObj.Type

	objTable := o.s.GetTable(table)
	if objTable == nil {
		return nil, errors.New("GetParentsViaChild: unknown object table " + table)
	}

	var parentObjs object.Array

	if objTable.ParentTables == nil {
		return nil, errors.New("GetParentsViaChild: cannot retrieve parents for table " + table + ", schema ParentTables is nil")
	}
	for _, pt := range objTable.ParentTables {
		pkQueryVals, err := pkQueryValsFromKV(childObj, o.s, pt)
		if err != nil {
			return nil, err
		}
		objs, err := o.RetrieveObjects(ctx, pt, pkQueryVals)
		if err != nil {
			return nil, err
		}
		parentObjs = append(parentObjs, objs...)
	}

	return parentObjs, nil
}

// NOTE: For foreign key filling, we do not (yet?) check to see if there are conflicts
// with regards to the uniqueness of primary key names.

// RetrieveWithChildren function will fleshen an *entire* object structure, given some primary keys
func (o ORM) RetrieveWithChildren(ctx context.Context, table string, pkValues map[string]interface{}) (*object.Object, error) {
	objTable := o.s.GetTable(table)
	if objTable == nil {
		return nil, errors.New("RetrieveWithChildren: unknown object table " + table)
	}

	obj, err := o.RetrieveObject(ctx, table, pkValues)
	if err != nil {
		return nil, errors.Wrap(err, "RetrieveWithChildren/RetrieveObject")
	}

	for name := range objTable.Children {
		childPkValues := make(map[string]interface{})

		childSchemaTable := o.s.Tables[name]
		pVal, ok := pkValues[childSchemaTable.Primary]
		if ok {
			childPkValues[childSchemaTable.Primary] = pVal
		}

		if childSchemaTable.MultiKey && childSchemaTable.ForeignKeys != nil {
			for _, fk := range childSchemaTable.ForeignKeys {
				childPkValues[fk] = pkValues[fk]
			}
		}
		childObj, err := o.RetrieveObject(ctx, name, childPkValues)
		if err != nil {
			return nil, errors.Wrap(err, "RetrieveWithChildren/RetrieveObject("+name+")")
		}
		if obj.Children[name] == nil {
			obj.Children[name] = make(object.Array, 1)
		}
		obj.Children[name][0] = childObj
	}
	return obj, nil
}

// RetrieveObject function will fleshen an object structure, given some primary keys
func (o ORM) RetrieveObject(ctx context.Context, table string, queryVals map[string]interface{}) (*object.Object, error) {
	// TODO: Implement LIMIT... That's all a singular retrieve should be underneath the hood that's different.
	objAry, err := o.RetrieveObjects(ctx, table, queryVals)
	if err != nil {
		return nil, err
	}
	if objAry == nil {
		return nil, nil
	}
	return objAry[0], nil
}

// FleshenChildren function accepts an object and resets it's children.
func (o ORM) FleshenChildren(ctx context.Context, obj *object.Object) (*object.Object, error) {
	schemaTable := o.s.GetTable(obj.Type)

	pkKey := schemaTable.Primary
	pkVal := obj.Get(pkKey)

	// If this table is configured with child tables, then we iterate over them and
	// call RetrieveObjects using the singular primary key value.
	// FIXME: We need to support multikey in this instance if we are going to consider
	// this complete.
	if len(schemaTable.Children) > 0 {
		for childTableName := range schemaTable.Children {
			m := map[string]interface{}{}
			m[pkKey] = pkVal
			childObjs, err := o.RetrieveObjects(ctx, childTableName, m)
			if err != nil {
				return nil, err
			}
			obj.Children[childTableName] = childObjs
		}
	}
	return obj, nil
}

// RetrieveObjectsFromCustomSQL will fleshen an object structure, given a SQL string.
func (o ORM) RetrieveObjectsFromCustomSQL(ctx context.Context, table string, sqlStr string, columnNames []string, bindArgs []interface{}) (object.Array, error) {
	if o.s == nil {
		return nil, errors.New("RetrieveObjectsFromCustomSQL: ORM schema set to nil?")
	}
	if o.s.Tables == nil {
		return nil, errors.New("RetrieveObjectsFromCustomSQL: ORM schema.Tables is set to nil?")
	}
	objTable := o.s.GetTable(table)
	if objTable == nil {
		return nil, errors.New("RetrieveObjectsFromCustomSQL: unknown object table " + table)
	}
	var objectArray object.Array

	if os.Getenv("DEBUG") != "" {
		fmt.Println("RetrieveObjectsFromCustomSQL/sqlStr=", sqlStr, "columnNames=", columnNames, "bindArgs=", bindArgs)
	}

	/*
		if err != nil {
			return nil, err
		}
	*/
	stmt, err := o.RawConn.PrepareContext(ctx, sqlStr)
	if err != nil {
		return nil, err
	}
	// TODO: Error not checked.
	defer stmt.Close()

	res, err := stmt.QueryContext(ctx, bindArgs...)
	if err != nil {
		return nil, err
	}
	// TODO: Error not checked.
	defer res.Close()
	columnTypes, err := res.ColumnTypes()
	if err != nil {
		return nil, err
	}

	for res.Next() {
		columnPointers, err := o.makeColumnPointers(len(columnNames), columnTypes)
		if err != nil {
			return nil, err
		}

		obj := object.New(table)
		if err := res.Scan(columnPointers...); err != nil {
			return nil, err
		}

		err = o.dynamicObjectSetter(columnNames, columnPointers, columnTypes, obj)
		if err != nil {
			return nil, err
		}
		obj.SetSaved(true)
		obj.ResetChangedFields()

		objectArray = append(objectArray, obj)
	}

	err = res.Err()
	if err != nil {
		return nil, err
	}
	return objectArray, nil
}

func (o ORM) makeQueryObj(objTable *schema.Table, queryVals map[string]interface{}) *object.Object {
	queryObj := object.New(objTable.Name)

	if objTable.FieldAliases == nil {
		queryObj.KV = queryVals
		return queryObj
	}
	for k, v := range queryVals {
		realName := objTable.GetFieldName(k)
		queryObj.KV[realName] = v
	}
	return queryObj
}

// RetrieveObjects function will fleshen an object structure, given some primary keys
func (o ORM) RetrieveObjects(ctx context.Context, table string, queryVals map[string]interface{}) (object.Array, error) {
	if o.s == nil {
		return nil, errors.New("RetrieveObjects: ORM schema set to nil?")
	}
	if o.s.Tables == nil {
		return nil, errors.New("RetrieveObjects: ORM schema.Tables is set to nil?")
	}
	objTable := o.s.GetTable(table)
	if objTable == nil {
		return nil, errors.New("RetrieveObjects: unknown object table " + table)
	}
	var objectArray object.Array

	queryObj := o.makeQueryObj(objTable, queryVals)

	sqlStr, columnNames, bindArgs, err := o.sqlGen.BindingRetrieve(o.s, queryObj)
	if os.Getenv("DEBUG") != "" {
		fmt.Println("RetrieveObjects/sqlStr=", sqlStr, "columnNames=", columnNames, "bindArgs=", bindArgs)
	}

	if err != nil {
		return nil, err
	}

	stmt, err := o.RawConn.PrepareContext(ctx, sqlStr)
	if err != nil {
		return nil, err
	}
	// TODO: Error not checked.
	defer stmt.Close()

	res, err := stmt.QueryContext(ctx, bindArgs...)
	if err != nil {
		return nil, err
	}
	// TODO: Error not checked.
	defer res.Close()
	columnTypes, err := res.ColumnTypes()
	if err != nil {
		return nil, err
	}

	for res.Next() {
		columnPointers, err := o.makeColumnPointers(len(columnNames), columnTypes)
		if err != nil {
			return nil, err
		}

		obj := object.New(table)
		if err := res.Scan(columnPointers...); err != nil {
			return nil, err
		}

		err = o.dynamicObjectSetter(columnNames, columnPointers, columnTypes, obj)
		if err != nil {
			return nil, err
		}
		obj.SetSaved(true)
		obj.ResetChangedFields()

		objectArray = append(objectArray, obj)
	}

	err = res.Err()
	if err != nil {
		return nil, err
	}
	return objectArray, nil
}

// NOTE: Read this post for more info on why the code below is written this way:
// https://stackoverflow.com/questions/23507531/is-golangs-sql-package-incapable-of-ad-hoc-exploratory-queries/23507765#23507765

func (o ORM) dynamicObjectSetter(columnNames []string, columnPointers []interface{}, columnTypes []*sql.ColumnType, obj *object.Object) error {
	sqlGen := o.sqlGen
	for i, v := range columnPointers {
		ct := columnTypes[i]

		typeName := ct.DatabaseTypeName()
		if sqlGen.IsStringType(typeName) || sqlGen.IsTimestampType(typeName) {
			nullable, _ := ct.Nullable()
			if nullable {
				val := v.(*sql.NullString)
				if val.Valid {
					obj.Set(columnNames[i], val.String)
				}
				// TODO: We don't set keys for null values. How else can we support this?
			} else {
				val := v.(*string)
				obj.Set(columnNames[i], *val)

			}
		} else if sqlGen.IsNumberType(typeName) {
			// TODO: support more than 'int64' for integer...
			nullable, _ := ct.Nullable()
			if nullable {
				val := v.(*sql.NullInt64)
				if val.Valid {
					obj.Set(columnNames[i], val.Int64)
				}
				// TODO: We don't set keys for null values. How else can we support this?
			} else {
				val := v.(*int64)
				obj.Set(columnNames[i], *val)
			}
		} else {
			return errors.New("dynamicObjectSetter: Unrecognized type: " + typeName)
		}
		// TODO: add timestamp support.
	}
	return nil
}

func (o ORM) makeColumnPointers(sliceLen int, columnTypes []*sql.ColumnType) ([]interface{}, error) {
	columnPointers := make([]interface{}, sliceLen)
	sqlGen := o.sqlGen
	for i := 0; i < sliceLen; i++ {
		ct := columnTypes[i]
		typeName := ct.DatabaseTypeName()
		if sqlGen.IsStringType(typeName) {
			nullable, _ := ct.Nullable()
			if nullable {
				var s sql.NullString
				columnPointers[i] = &s
			} else {
				var s string
				columnPointers[i] = &s
			}
		} else if sqlGen.IsNumberType(typeName) {
			nullable, _ := ct.Nullable()
			if nullable {
				var j sql.NullInt64
				columnPointers[i] = &j
			} else {
				var j int64
				columnPointers[i] = &j

			}
		} else if sqlGen.IsTimestampType(typeName) {
			nullable, _ := ct.Nullable()
			if nullable {
				var j sql.NullString
				columnPointers[i] = &j
			} else {
				var j string
				columnPointers[i] = &j

			}
		} else {
			return nil, errors.New("makeColumnPointers: Unrecognized type: " + typeName)
		}
	}
	return columnPointers, nil
}
<commit_msg>minor changes, debugging<commit_after>package orm

import (
	"context"
	"database/sql"
	"fmt"
	"os"

	"github.com/pkg/errors"
	"github.com/rbastic/dyndao/object"
	"github.com/rbastic/dyndao/schema"
)

// GetParentsViaChild retrieves all direct (one-level 'up') parents for a given child object.
// If a child contains multiple parent tables (possibility?) then this would return an Array
// of objects with multiple potential values for their obj.Type fields.
func (o ORM) GetParentsViaChild(ctx context.Context, childObj *object.Object) (object.Array, error) {
	table := childObj.Type

	objTable := o.s.GetTable(table)
	if objTable == nil {
		return nil, errors.New("GetParentsViaChild: unknown object table " + table)
	}

	var parentObjs object.Array

	if objTable.ParentTables == nil {
		return nil, errors.New("GetParentsViaChild: cannot retrieve parents for table " + table + ", schema ParentTables is nil")
	}
	for _, pt := range objTable.ParentTables {
		pkQueryVals, err := pkQueryValsFromKV(childObj, o.s, pt)
		if err != nil {
			return nil, err
		}
		objs, err := o.RetrieveObjects(ctx, pt, pkQueryVals)
		if err != nil {
			return nil, err
		}
		parentObjs = append(parentObjs, objs...)
	}

	return parentObjs, nil
}

// NOTE: For foreign key filling, we do not (yet?) check to see if there are conflicts
// with regards to the uniqueness of primary key names.

// RetrieveWithChildren function will fleshen an *entire* object structure, given some primary keys
func (o ORM) RetrieveWithChildren(ctx context.Context, table string, pkValues map[string]interface{}) (*object.Object, error) {
	objTable := o.s.GetTable(table)
	if objTable == nil {
		return nil, errors.New("RetrieveWithChildren: unknown object table " + table)
	}

	obj, err := o.RetrieveObject(ctx, table, pkValues)
	if err != nil {
		return nil, errors.Wrap(err, "RetrieveWithChildren/RetrieveObject")
	}

	for name := range objTable.Children {
		childPkValues := make(map[string]interface{})

		childSchemaTable := o.s.Tables[name]
		pVal, ok := pkValues[childSchemaTable.Primary]
		if ok {
			childPkValues[childSchemaTable.Primary] = pVal
		}

		if childSchemaTable.MultiKey && childSchemaTable.ForeignKeys != nil {
			for _, fk := range childSchemaTable.ForeignKeys {
				childPkValues[fk] = pkValues[fk]
			}
		}
		childObj, err := o.RetrieveObject(ctx, name, childPkValues)
		if err != nil {
			return nil, errors.Wrap(err, "RetrieveWithChildren/RetrieveObject("+name+")")
		}
		if obj.Children[name] == nil {
			obj.Children[name] = make(object.Array, 1)
		}
		obj.Children[name][0] = childObj
	}
	return obj, nil
}

// RetrieveObject function will fleshen an object structure, given some primary keys
func (o ORM) RetrieveObject(ctx context.Context, table string, queryVals map[string]interface{}) (*object.Object, error) {
	// TODO: Implement LIMIT... That's all a singular retrieve should be underneath the hood that's different.
	objAry, err := o.RetrieveObjects(ctx, table, queryVals)
	if err != nil {
		return nil, err
	}
	if objAry == nil {
		return nil, nil
	}
	return objAry[0], nil
}

// FleshenChildren function accepts an object and resets it's children.
func (o ORM) FleshenChildren(ctx context.Context, obj *object.Object) (*object.Object, error) {
	schemaTable := o.s.GetTable(obj.Type)

	pkKey := schemaTable.Primary
	pkVal := obj.Get(pkKey)

	// If this table is configured with child tables, then we iterate over
	// them and call RetrieveObjects using the singular primary key value.
	// FIXME: We need to support multikey in this instance if we are going
	// to consider this complete.
	if len(schemaTable.Children) > 0 {
		for childTableName := range schemaTable.Children {
			m := map[string]interface{}{}
			m[pkKey] = pkVal
			childObjs, err := o.RetrieveObjects(ctx, childTableName, m)
			if err != nil {
				return nil, err
			}
			obj.Children[childTableName] = childObjs
		}
	}
	return obj, nil
}

// RetrieveObjectsFromCustomSQL will fleshen an object structure, given a SQL string.
func (o ORM) RetrieveObjectsFromCustomSQL(ctx context.Context, table string, sqlStr string, columnNames []string, bindArgs []interface{}) (object.Array, error) {
	objTable := o.s.GetTable(table)
	if objTable == nil {
		return nil, errors.New("RetrieveObjectsFromCustomSQL: unknown object table " + table)
	}
	var objectArray object.Array

	if os.Getenv("DEBUG") != "" {
		fmt.Println("RetrieveObjectsFromCustomSQL/sqlStr=", sqlStr, "columnNames=", columnNames, "bindArgs=", bindArgs)
	}

	stmt, err := o.RawConn.PrepareContext(ctx, sqlStr)
	if err != nil {
		return nil, err
	}
	// TODO: Error not checked.
	defer stmt.Close()

	res, err := stmt.QueryContext(ctx, bindArgs...)
	if err != nil {
		return nil, err
	}
	// TODO: Error not checked.
	defer res.Close()
	columnTypes, err := res.ColumnTypes()
	if err != nil {
		return nil, err
	}

	for res.Next() {
		columnPointers, err := o.makeColumnPointers(len(columnNames), columnTypes)
		if err != nil {
			return nil, err
		}

		obj := object.New(table)
		if err := res.Scan(columnPointers...); err != nil {
			return nil, err
		}

		err = o.dynamicObjectSetter(columnNames, columnPointers, columnTypes, obj)
		if err != nil {
			return nil, err
		}
		obj.SetSaved(true)
		obj.ResetChangedFields()

		objectArray = append(objectArray, obj)
	}

	err = res.Err()
	if err != nil {
		return nil, err
	}
	return objectArray, nil
}

func (o ORM) makeQueryObj(objTable *schema.Table, queryVals map[string]interface{}) *object.Object {
	queryObj := object.New(objTable.Name)

	if objTable.FieldAliases == nil {
		queryObj.KV = queryVals
		return queryObj
	}
	for k, v := range queryVals {
		realName := objTable.GetFieldName(k)
		queryObj.KV[realName] = v
	}
	return queryObj
}

// RetrieveObjects function will fleshen an object structure, given some primary keys
func (o ORM) RetrieveObjects(ctx context.Context, table string, queryVals map[string]interface{}) (object.Array, error) {
	objTable := o.s.GetTable(table)
	if objTable == nil {
		return nil, errors.New("RetrieveObjects: unknown object table " + table)
	}
	var objectArray object.Array
	queryObj := o.makeQueryObj(objTable, queryVals)

	sqlStr, columnNames, bindArgs, err := o.sqlGen.BindingRetrieve(o.s, queryObj)
	if os.Getenv("DEBUG") != "" {
		fmt.Println("RetrieveObjects/sqlStr=", sqlStr, "columnNames=", columnNames, "bindArgs=", bindArgs)
	}

	if err != nil {
		return nil, err
	}

	stmt, err := o.RawConn.PrepareContext(ctx, sqlStr)
	if err != nil {
		return nil, err
	}
	// TODO: Error not checked.
	defer stmt.Close()

	res, err := stmt.QueryContext(ctx, bindArgs...)
	if err != nil {
		return nil, err
	}
	// TODO: Error not checked.
	defer res.Close()
	columnTypes, err := res.ColumnTypes()
	if err != nil {
		return nil, err
	}

	for res.Next() {
		columnPointers, err := o.makeColumnPointers(len(columnNames), columnTypes)
		if err != nil {
			return nil, err
		}

		obj := object.New(table)
		if err := res.Scan(columnPointers...); err != nil {
			return nil, err
		}

		err = o.dynamicObjectSetter(columnNames, columnPointers, columnTypes, obj)
		if err != nil {
			return nil, err
		}
		obj.SetSaved(true)
		obj.ResetChangedFields()

		if os.Getenv("DEBUG") != "" {
			fmt.Println("RetrieveObjects ... obj ->", obj)
		}
		objectArray = append(objectArray, obj)
	}

	err = res.Err()
	if err != nil {
		return nil, err
	}
	return objectArray, nil
}

// NOTE: Read this post for more info on why the code below is written this way:
// https://stackoverflow.com/questions/23507531/is-golangs-sql-package-incapable-of-ad-hoc-exploratory-queries/23507765#23507765

func (o ORM) dynamicObjectSetter(columnNames []string, columnPointers []interface{}, columnTypes []*sql.ColumnType, obj *object.Object) error {
	sqlGen := o.sqlGen
	for i, v := range columnPointers {
		ct := columnTypes[i]

		typeName := ct.DatabaseTypeName()
		if sqlGen.IsStringType(typeName) || sqlGen.IsTimestampType(typeName) {
			nullable, _ := ct.Nullable()
			if nullable {
				val := v.(*sql.NullString)
				if val.Valid {
					obj.Set(columnNames[i], val.String)
				}
				// TODO: We don't set keys for null values. How else can we support this?
			} else {
				val := v.(*string)
				obj.Set(columnNames[i], *val)

			}
		} else if sqlGen.IsNumberType(typeName) {
			// TODO: support more than 'int64' for integer...
			nullable, _ := ct.Nullable()
			if nullable {
				val := v.(*sql.NullInt64)
				if val.Valid {
					obj.Set(columnNames[i], val.Int64)
				}
				// TODO: We don't set keys for null values. How else can we support this?
			} else {
				val := v.(*int64)
				obj.Set(columnNames[i], *val)
			}
		} else {
			return errors.New("dynamicObjectSetter: Unrecognized type: " + typeName)
		}
		// TODO: add timestamp support.
	}
	return nil
}

func (o ORM) makeColumnPointers(sliceLen int, columnTypes []*sql.ColumnType) ([]interface{}, error) {
	columnPointers := make([]interface{}, sliceLen)
	sqlGen := o.sqlGen
	for i := 0; i < sliceLen; i++ {
		ct := columnTypes[i]
		typeName := ct.DatabaseTypeName()
		if sqlGen.IsStringType(typeName) {
			nullable, _ := ct.Nullable()
			if nullable {
				var s sql.NullString
				columnPointers[i] = &s
			} else {
				var s string
				columnPointers[i] = &s
			}
		} else if sqlGen.IsNumberType(typeName) {
			nullable, _ := ct.Nullable()
			if nullable {
				var j sql.NullInt64
				columnPointers[i] = &j
			} else {
				var j int64
				columnPointers[i] = &j

			}
		} else if sqlGen.IsTimestampType(typeName) {
			nullable, _ := ct.Nullable()
			if nullable {
				var j sql.NullString
				columnPointers[i] = &j
			} else {
				var j string
				columnPointers[i] = &j

			}
		} else {
			return nil, errors.New("makeColumnPointers: Unrecognized type: " + typeName)
		}
	}
	return columnPointers, nil
}
<|endoftext|>


<commit_before>package restful

// Copyright 2013 Ernest Micklei. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

import (
	"bytes"
	"fmt"
	"net/http"
	"os"
	"runtime"
	"strings"
	"sync"

	"github.com/emicklei/go-restful/log"
)

// Container holds a collection of WebServices and a http.ServeMux to dispatch http requests.
// The requests are further dispatched to routes of WebServices using a RouteSelector
type Container struct {
	webServicesLock        sync.RWMutex
	webServices            []*WebService
	ServeMux               *http.ServeMux
	isRegisteredOnRoot     bool
	containerFilters       []FilterFunction
	doNotRecover           bool // default is false
	recoverHandleFunc      RecoverHandleFunction
	serviceErrorHandleFunc ServiceErrorHandleFunction
	router                 RouteSelector // default is a RouterJSR311, CurlyRouter is the faster alternative
	contentEncodingEnabled bool          // default is false
}

// NewContainer creates a new Container using a new ServeMux and default router (RouterJSR311)
func NewContainer() *Container {
	return &Container{
		webServices:            []*WebService{},
		ServeMux:               http.NewServeMux(),
		isRegisteredOnRoot:     false,
		containerFilters:       []FilterFunction{},
		doNotRecover:           false,
		recoverHandleFunc:      logStackOnRecover,
		serviceErrorHandleFunc: writeServiceError,
		router:                 RouterJSR311{},
		contentEncodingEnabled: false}
}

// RecoverHandleFunction declares functions that can be used to handle a panic situation.
// The first argument is what recover() returns. The second must be used to communicate an error response.
type RecoverHandleFunction func(interface{}, http.ResponseWriter)

// RecoverHandler changes the default function (logStackOnRecover) to be called
// when a panic is detected. DoNotRecover must be have its default value (=false).
func (c *Container) RecoverHandler(handler RecoverHandleFunction) {
	c.recoverHandleFunc = handler
}

// ServiceErrorHandleFunction declares functions that can be used to handle a service error situation.
// The first argument is the service error, the second is the request that resulted in the error and
// the third must be used to communicate an error response.
type ServiceErrorHandleFunction func(ServiceError, *Request, *Response)

// ServiceErrorHandler changes the default function (writeServiceError) to be called
// when a ServiceError is detected.
func (c *Container) ServiceErrorHandler(handler ServiceErrorHandleFunction) {
	c.serviceErrorHandleFunc = handler
}

// DoNotRecover controls whether panics will be caught to return HTTP 500.
// If set to true, Route functions are responsible for handling any error situation.
// Default value is false = recover from panics. This has performance implications.
func (c *Container) DoNotRecover(doNot bool) {
	c.doNotRecover = doNot
}

// Router changes the default Router (currently RouterJSR311)
func (c *Container) Router(aRouter RouteSelector) {
	c.router = aRouter
}

// EnableContentEncoding (default=false) allows for GZIP or DEFLATE encoding of responses.
func (c *Container) EnableContentEncoding(enabled bool) {
	c.contentEncodingEnabled = enabled
}

// Add a WebService to the Container. It will detect duplicate root paths and panic in that case.
func (c *Container) Add(service *WebService) *Container {
	c.webServicesLock.Lock()
	defer c.webServicesLock.Unlock()
	// If registered on root then no additional specific mapping is needed
	if !c.isRegisteredOnRoot {
		pattern := c.fixedPrefixPath(service.RootPath())
		// check if root path registration is needed
		if "/" == pattern || "" == pattern {
			c.ServeMux.HandleFunc("/", c.dispatch)
			c.isRegisteredOnRoot = true
		} else {
			// detect if registration already exists
			alreadyMapped := false
			for _, each := range c.webServices {
				if each.RootPath() == service.RootPath() {
					alreadyMapped = true
					break
				}
			}
			if !alreadyMapped {
				c.ServeMux.HandleFunc(pattern, c.dispatch)
				if !strings.HasSuffix(pattern, "/") {
					c.ServeMux.HandleFunc(pattern+"/", c.dispatch)
				}
			}
		}
	}
	// cannot have duplicate root paths
	for _, each := range c.webServices {
		if each.RootPath() == service.RootPath() {
			log.Printf("[restful] WebService with duplicate root path detected:['%v']", each)
			os.Exit(1)
		}
	}
	// if rootPath was not set then lazy initialize it
	if len(service.rootPath) == 0 {
		service.Path("/")
	}
	c.webServices = append(c.webServices, service)
	return c
}

func (c *Container) Remove(ws *WebService) error {
	c.webServicesLock.Lock()
	defer c.webServicesLock.Unlock()
	newServices := []*WebService{}
	for ix := range c.webServices {
		if c.webServices[ix].rootPath != ws.rootPath {
			newServices = append(newServices, c.webServices[ix])
		}
	}
	c.webServices = newServices
	return nil
}

// logStackOnRecover is the default RecoverHandleFunction and is called
// when DoNotRecover is false and the recoverHandleFunc is not set for the container.
// Default implementation logs the stacktrace and writes the stacktrace on the response.
// This may be a security issue as it exposes sourcecode information.
func logStackOnRecover(panicReason interface{}, httpWriter http.ResponseWriter) {
	var buffer bytes.Buffer
	buffer.WriteString(fmt.Sprintf("[restful] recover from panic situation: - %v\r\n", panicReason))
	for i := 2; ; i += 1 {
		_, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		buffer.WriteString(fmt.Sprintf("    %s:%d\r\n", file, line))
	}
	log.Print(buffer.String())
	httpWriter.WriteHeader(http.StatusInternalServerError)
	httpWriter.Write(buffer.Bytes())
}

// writeServiceError is the default ServiceErrorHandleFunction and is called
// when a ServiceError is returned during route selection. Default implementation
// calls resp.WriteErrorString(err.Code, err.Message)
func writeServiceError(err ServiceError, req *Request, resp *Response) {
	resp.WriteErrorString(err.Code, err.Message)
}

// Dispatch the incoming Http Request to a matching WebService.
func (c *Container) dispatch(httpWriter http.ResponseWriter, httpRequest *http.Request) {
	writer := httpWriter

	// CompressingResponseWriter should be closed after all operations are done
	defer func() {
		if compressWriter, ok := writer.(*CompressingResponseWriter); ok {
			compressWriter.Close()
		}
	}()

	// Instal panic recovery unless told otherwise
	if !c.doNotRecover { // catch all for 500 response
		defer func() {
			if r := recover(); r != nil {
				c.recoverHandleFunc(r, writer)
				return
			}
		}()
	}
	// Install closing the request body (if any)
	defer func() {
		if nil != httpRequest.Body {
			httpRequest.Body.Close()
		}
	}()

	// Detect if compression is needed
	// assume without compression, test for override
	if c.contentEncodingEnabled {
		doCompress, encoding := wantsCompressedResponse(httpRequest)
		if doCompress {
			var err error
			writer, err = NewCompressingResponseWriter(httpWriter, encoding)
			if err != nil {
				log.Print("[restful] unable to install compressor: ", err)
				httpWriter.WriteHeader(http.StatusInternalServerError)
				return
			}
		}
	}
	// Find best match Route ; err is non nil if no match was found
	var webService *WebService
	var route *Route
	var err error
	func() {
		c.webServicesLock.RLock()
		defer c.webServicesLock.RUnlock()
		webService, route, err = c.router.SelectRoute(
			c.webServices,
			httpRequest)
	}()
	if err != nil {
		// a non-200 response has already been written
		// run container filters anyway ; they should not touch the response...
		chain := FilterChain{Filters: c.containerFilters, Target: func(req *Request, resp *Response) {
			switch err.(type) {
			case ServiceError:
				ser := err.(ServiceError)
				c.serviceErrorHandleFunc(ser, req, resp)
			}
			// TODO
		}}
		chain.ProcessFilter(NewRequest(httpRequest), NewResponse(writer))
		return
	}
	wrappedRequest, wrappedResponse := route.wrapRequestResponse(writer, httpRequest)
	// pass through filters (if any)
	if len(c.containerFilters)+len(webService.filters)+len(route.Filters) > 0 {
		// compose filter chain
		allFilters := []FilterFunction{}
		allFilters = append(allFilters, c.containerFilters...)
		allFilters = append(allFilters, webService.filters...)
		allFilters = append(allFilters, route.Filters...)
		chain := FilterChain{Filters: allFilters, Target: func(req *Request, resp *Response) {
			// handle request by route after passing all filters
			route.Function(wrappedRequest, wrappedResponse)
		}}
		chain.ProcessFilter(wrappedRequest, wrappedResponse)
	} else {
		// no filters, handle request by route
		route.Function(wrappedRequest, wrappedResponse)
	}
}

// fixedPrefixPath returns the fixed part of the partspec ; it may include template vars {}
func (c Container) fixedPrefixPath(pathspec string) string {
	varBegin := strings.Index(pathspec, "{")
	if -1 == varBegin {
		return pathspec
	}
	return pathspec[:varBegin]
}

// ServeHTTP implements net/http.Handler therefore a Container can be a Handler in a http.Server
func (c Container) ServeHTTP(httpwriter http.ResponseWriter, httpRequest *http.Request) {
	c.ServeMux.ServeHTTP(httpwriter, httpRequest)
}

// Handle registers the handler for the given pattern. If a handler already exists for pattern, Handle panics.
func (c Container) Handle(pattern string, handler http.Handler) {
	c.ServeMux.Handle(pattern, handler)
}

// HandleWithFilter registers the handler for the given pattern.
// Container's filter chain is applied for handler.
// If a handler already exists for pattern, HandleWithFilter panics.
func (c *Container) HandleWithFilter(pattern string, handler http.Handler) {
	f := func(httpResponse http.ResponseWriter, httpRequest *http.Request) {
		if len(c.containerFilters) == 0 {
			handler.ServeHTTP(httpResponse, httpRequest)
			return
		}

		chain := FilterChain{Filters: c.containerFilters, Target: func(req *Request, resp *Response) {
			handler.ServeHTTP(httpResponse, httpRequest)
		}}
		chain.ProcessFilter(NewRequest(httpRequest), NewResponse(httpResponse))
	}

	c.Handle(pattern, http.HandlerFunc(f))
}

// Filter appends a container FilterFunction. These are called before dispatching
// a http.Request to a WebService from the container
func (c *Container) Filter(filter FilterFunction) {
	c.containerFilters = append(c.containerFilters, filter)
}

// RegisteredWebServices returns the collections of added WebServices
func (c Container) RegisteredWebServices() []*WebService {
	c.webServicesLock.RLock()
	defer c.webServicesLock.RUnlock()
	result := make([]*WebService, len(c.webServices))
	for ix := range c.webServices {
		result[ix] = c.webServices[ix]
	}
	return result
}

// computeAllowedMethods returns a list of HTTP methods that are valid for a Request
func (c Container) computeAllowedMethods(req *Request) []string {
	// Go through all RegisteredWebServices() and all its Routes to collect the options
	methods := []string{}
	requestPath := req.Request.URL.Path
	for _, ws := range c.RegisteredWebServices() {
		matches := ws.pathExpr.Matcher.FindStringSubmatch(requestPath)
		if matches != nil {
			finalMatch := matches[len(matches)-1]
			for _, rt := range ws.Routes() {
				matches := rt.pathExpr.Matcher.FindStringSubmatch(finalMatch)
				if matches != nil {
					lastMatch := matches[len(matches)-1]
					if lastMatch == "" || lastMatch == "/" { // do not include if value is neither empty nor ‘/’.
						methods = append(methods, rt.Method)
					}
				}
			}
		}
	}
	// methods = append(methods, "OPTIONS")  not sure about this
	return methods
}

// newBasicRequestResponse creates a pair of Request,Response from its http versions.
// It is basic because no parameter or (produces) content-type information is given.
func newBasicRequestResponse(httpWriter http.ResponseWriter, httpRequest *http.Request) (*Request, *Response) {
	resp := NewResponse(httpWriter)
	resp.requestAccept = httpRequest.Header.Get(HEADER_Accept)
	return NewRequest(httpRequest), resp
}
<commit_msg>Don't re-register dispatch handlers<commit_after>package restful

// Copyright 2013 Ernest Micklei. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

import (
	"bytes"
	"fmt"
	"net/http"
	"os"
	"runtime"
	"strings"
	"sync"

	"github.com/emicklei/go-restful/log"
)

// Container holds a collection of WebServices and a http.ServeMux to dispatch http requests.
// The requests are further dispatched to routes of WebServices using a RouteSelector
type Container struct {
	webServicesLock        sync.RWMutex
	webServices            []*WebService
	ServeMux               *http.ServeMux
	isRegisteredOnRoot     bool
	containerFilters       []FilterFunction
	doNotRecover           bool // default is false
	recoverHandleFunc      RecoverHandleFunction
	serviceErrorHandleFunc ServiceErrorHandleFunction
	router                 RouteSelector // default is a RouterJSR311, CurlyRouter is the faster alternative
	contentEncodingEnabled bool          // default is false
	registrations          map[string]bool
}

// NewContainer creates a new Container using a new ServeMux and default router (RouterJSR311)
func NewContainer() *Container {
	return &Container{
		webServices:            []*WebService{},
		ServeMux:               http.NewServeMux(),
		isRegisteredOnRoot:     false,
		containerFilters:       []FilterFunction{},
		doNotRecover:           false,
		recoverHandleFunc:      logStackOnRecover,
		serviceErrorHandleFunc: writeServiceError,
		router:                 RouterJSR311{},
		contentEncodingEnabled: false,
		registrations:          map[string]bool{},
	}
}

// RecoverHandleFunction declares functions that can be used to handle a panic situation.
// The first argument is what recover() returns. The second must be used to communicate an error response.
type RecoverHandleFunction func(interface{}, http.ResponseWriter)

// RecoverHandler changes the default function (logStackOnRecover) to be called
// when a panic is detected. DoNotRecover must be have its default value (=false).
func (c *Container) RecoverHandler(handler RecoverHandleFunction) {
	c.recoverHandleFunc = handler
}

// ServiceErrorHandleFunction declares functions that can be used to handle a service error situation.
// The first argument is the service error, the second is the request that resulted in the error and
// the third must be used to communicate an error response.
type ServiceErrorHandleFunction func(ServiceError, *Request, *Response)

// ServiceErrorHandler changes the default function (writeServiceError) to be called
// when a ServiceError is detected.
func (c *Container) ServiceErrorHandler(handler ServiceErrorHandleFunction) {
	c.serviceErrorHandleFunc = handler
}

// DoNotRecover controls whether panics will be caught to return HTTP 500.
// If set to true, Route functions are responsible for handling any error situation.
// Default value is false = recover from panics. This has performance implications.
func (c *Container) DoNotRecover(doNot bool) {
	c.doNotRecover = doNot
}

// Router changes the default Router (currently RouterJSR311)
func (c *Container) Router(aRouter RouteSelector) {
	c.router = aRouter
}

// EnableContentEncoding (default=false) allows for GZIP or DEFLATE encoding of responses.
func (c *Container) EnableContentEncoding(enabled bool) {
	c.contentEncodingEnabled = enabled
}

// Add a WebService to the Container. It will detect duplicate root paths and panic in that case.
func (c *Container) Add(service *WebService) *Container {
	c.webServicesLock.Lock()
	defer c.webServicesLock.Unlock()
	// If registered on root then no additional specific mapping is needed
	if !c.isRegisteredOnRoot {
		pattern := c.fixedPrefixPath(service.RootPath())
		// check if root path registration is needed
		if "/" == pattern || "" == pattern {
			c.HandleFuncDispatch("/")
			c.isRegisteredOnRoot = true
		} else {
			// detect if registration already exists
			alreadyMapped := false
			for _, each := range c.webServices {
				if each.RootPath() == service.RootPath() {
					alreadyMapped = true
					break
				}
			}
			if !alreadyMapped {
				c.HandleFuncDispatch(pattern)
				if !strings.HasSuffix(pattern, "/") {
					c.HandleFuncDispatch(pattern + "/")
				}
			}
		}
	}
	// cannot have duplicate root paths
	for _, each := range c.webServices {
		if each.RootPath() == service.RootPath() {
			log.Printf("[restful] WebService with duplicate root path detected:['%v']", each)
			os.Exit(1)
		}
	}
	// if rootPath was not set then lazy initialize it
	if len(service.rootPath) == 0 {
		service.Path("/")
	}
	c.webServices = append(c.webServices, service)
	return c
}

func (c *Container) Remove(ws *WebService) error {
	c.webServicesLock.Lock()
	defer c.webServicesLock.Unlock()
	newServices := []*WebService{}
	for ix := range c.webServices {
		if c.webServices[ix].rootPath != ws.rootPath {
			newServices = append(newServices, c.webServices[ix])
		}
	}
	c.webServices = newServices
	return nil
}

// logStackOnRecover is the default RecoverHandleFunction and is called
// when DoNotRecover is false and the recoverHandleFunc is not set for the container.
// Default implementation logs the stacktrace and writes the stacktrace on the response.
// This may be a security issue as it exposes sourcecode information.
func logStackOnRecover(panicReason interface{}, httpWriter http.ResponseWriter) {
	var buffer bytes.Buffer
	buffer.WriteString(fmt.Sprintf("[restful] recover from panic situation: - %v\r\n", panicReason))
	for i := 2; ; i += 1 {
		_, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		buffer.WriteString(fmt.Sprintf("    %s:%d\r\n", file, line))
	}
	log.Print(buffer.String())
	httpWriter.WriteHeader(http.StatusInternalServerError)
	httpWriter.Write(buffer.Bytes())
}

// writeServiceError is the default ServiceErrorHandleFunction and is called
// when a ServiceError is returned during route selection. Default implementation
// calls resp.WriteErrorString(err.Code, err.Message)
func writeServiceError(err ServiceError, req *Request, resp *Response) {
	resp.WriteErrorString(err.Code, err.Message)
}

// Dispatch the incoming Http Request to a matching WebService.
func (c *Container) dispatch(httpWriter http.ResponseWriter, httpRequest *http.Request) {
	writer := httpWriter

	// CompressingResponseWriter should be closed after all operations are done
	defer func() {
		if compressWriter, ok := writer.(*CompressingResponseWriter); ok {
			compressWriter.Close()
		}
	}()

	// Instal panic recovery unless told otherwise
	if !c.doNotRecover { // catch all for 500 response
		defer func() {
			if r := recover(); r != nil {
				c.recoverHandleFunc(r, writer)
				return
			}
		}()
	}
	// Install closing the request body (if any)
	defer func() {
		if nil != httpRequest.Body {
			httpRequest.Body.Close()
		}
	}()

	// Detect if compression is needed
	// assume without compression, test for override
	if c.contentEncodingEnabled {
		doCompress, encoding := wantsCompressedResponse(httpRequest)
		if doCompress {
			var err error
			writer, err = NewCompressingResponseWriter(httpWriter, encoding)
			if err != nil {
				log.Print("[restful] unable to install compressor: ", err)
				httpWriter.WriteHeader(http.StatusInternalServerError)
				return
			}
		}
	}
	// Find best match Route ; err is non nil if no match was found
	var webService *WebService
	var route *Route
	var err error
	func() {
		c.webServicesLock.RLock()
		defer c.webServicesLock.RUnlock()
		webService, route, err = c.router.SelectRoute(
			c.webServices,
			httpRequest)
	}()
	if err != nil {
		// a non-200 response has already been written
		// run container filters anyway ; they should not touch the response...
		chain := FilterChain{Filters: c.containerFilters, Target: func(req *Request, resp *Response) {
			switch err.(type) {
			case ServiceError:
				ser := err.(ServiceError)
				c.serviceErrorHandleFunc(ser, req, resp)
			}
			// TODO
		}}
		chain.ProcessFilter(NewRequest(httpRequest), NewResponse(writer))
		return
	}
	wrappedRequest, wrappedResponse := route.wrapRequestResponse(writer, httpRequest)
	// pass through filters (if any)
	if len(c.containerFilters)+len(webService.filters)+len(route.Filters) > 0 {
		// compose filter chain
		allFilters := []FilterFunction{}
		allFilters = append(allFilters, c.containerFilters...)
		allFilters = append(allFilters, webService.filters...)
		allFilters = append(allFilters, route.Filters...)
		chain := FilterChain{Filters: allFilters, Target: func(req *Request, resp *Response) {
			// handle request by route after passing all filters
			route.Function(wrappedRequest, wrappedResponse)
		}}
		chain.ProcessFilter(wrappedRequest, wrappedResponse)
	} else {
		// no filters, handle request by route
		route.Function(wrappedRequest, wrappedResponse)
	}
}

// fixedPrefixPath returns the fixed part of the partspec ; it may include template vars {}
func (c Container) fixedPrefixPath(pathspec string) string {
	varBegin := strings.Index(pathspec, "{")
	if -1 == varBegin {
		return pathspec
	}
	return pathspec[:varBegin]
}

// ServeHTTP implements net/http.Handler therefore a Container can be a Handler in a http.Server
func (c Container) ServeHTTP(httpwriter http.ResponseWriter, httpRequest *http.Request) {
	c.ServeMux.ServeHTTP(httpwriter, httpRequest)
}

// Handle registers the handler for the given pattern. If a handler already exists for pattern, Handle panics.
func (c Container) Handle(pattern string, handler http.Handler) {
	c.ServeMux.Handle(pattern, handler)
}

// HandleFuncDispatch register for the given pattern.
// Unlike the above, it first checks to see if that pattern is already
// registered and skips registration.  This is because once you start handling
// something in ServeMux you can't remove it.  (which seems broken imho)
// Requires webServicesLock to be held.
func (c *Container) HandleFuncDispatch(pattern string) {
	_, found := c.registrations[pattern]
	if !found {
		c.ServeMux.HandleFunc(pattern, c.dispatch)
		c.registrations[pattern] = true
	}
}

// HandleWithFilter registers the handler for the given pattern.
// Container's filter chain is applied for handler.
// If a handler already exists for pattern, HandleWithFilter panics.
func (c *Container) HandleWithFilter(pattern string, handler http.Handler) {
	f := func(httpResponse http.ResponseWriter, httpRequest *http.Request) {
		if len(c.containerFilters) == 0 {
			handler.ServeHTTP(httpResponse, httpRequest)
			return
		}

		chain := FilterChain{Filters: c.containerFilters, Target: func(req *Request, resp *Response) {
			handler.ServeHTTP(httpResponse, httpRequest)
		}}
		chain.ProcessFilter(NewRequest(httpRequest), NewResponse(httpResponse))
	}

	c.Handle(pattern, http.HandlerFunc(f))
}

// Filter appends a container FilterFunction. These are called before dispatching
// a http.Request to a WebService from the container
func (c *Container) Filter(filter FilterFunction) {
	c.containerFilters = append(c.containerFilters, filter)
}

// RegisteredWebServices returns the collections of added WebServices
func (c Container) RegisteredWebServices() []*WebService {
	c.webServicesLock.RLock()
	defer c.webServicesLock.RUnlock()
	result := make([]*WebService, len(c.webServices))
	for ix := range c.webServices {
		result[ix] = c.webServices[ix]
	}
	return result
}

// computeAllowedMethods returns a list of HTTP methods that are valid for a Request
func (c Container) computeAllowedMethods(req *Request) []string {
	// Go through all RegisteredWebServices() and all its Routes to collect the options
	methods := []string{}
	requestPath := req.Request.URL.Path
	for _, ws := range c.RegisteredWebServices() {
		matches := ws.pathExpr.Matcher.FindStringSubmatch(requestPath)
		if matches != nil {
			finalMatch := matches[len(matches)-1]
			for _, rt := range ws.Routes() {
				matches := rt.pathExpr.Matcher.FindStringSubmatch(finalMatch)
				if matches != nil {
					lastMatch := matches[len(matches)-1]
					if lastMatch == "" || lastMatch == "/" { // do not include if value is neither empty nor ‘/’.
						methods = append(methods, rt.Method)
					}
				}
			}
		}
	}
	// methods = append(methods, "OPTIONS")  not sure about this
	return methods
}

// newBasicRequestResponse creates a pair of Request,Response from its http versions.
// It is basic because no parameter or (produces) content-type information is given.
func newBasicRequestResponse(httpWriter http.ResponseWriter, httpRequest *http.Request) (*Request, *Response) {
	resp := NewResponse(httpWriter)
	resp.requestAccept = httpRequest.Header.Get(HEADER_Accept)
	return NewRequest(httpRequest), resp
}
<|endoftext|>


<commit_before>// Copyright ©2015 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package conv provides matrix type interconversion utilities.
package conv

import (
	"github.com/gonum/matrix"
	"github.com/gonum/matrix/cmat128"
	"github.com/gonum/matrix/mat64"
)

// Complex is a complex matrix constructed from two real matrices.
type Complex struct {
	// r and i are not exposed to ensure that
	// their dimensions can not be altered by
	// clients behind our back.
	r, i mat64.Matrix
}

var (
	_ Realer = Complex{}
	_ Imager = Complex{}
)

// NewComplex returns a complex matrix constructed from r and i. At least one of
// r or i must be non-nil otherwise NewComples will panic. If one of the inputs
// is nil, that part of the complex number will be zero when returned by At.
// If both are non-nil but differ in their sizes, NewComplex will panic.
func NewComplex(r, i mat64.Matrix) Complex {
	if r == nil && i == nil {
		panic("conv: no matrix")
	} else if r != nil && i != nil {
		rr, rc := r.Dims()
		ir, ic := i.Dims()
		if rr != ir || rc != ic {
			panic(matrix.ErrShape)
		}
	}
	return Complex{r: r, i: i}
}

// Dims returns the number of rows and columns in the matrix.
func (m Complex) Dims() (r, c int) {
	if m.r == nil {
		return m.i.Dims()
	}
	return m.r.Dims()
}

// At returns the element at row r, column c.
func (m Complex) At(r, c int) complex128 {
	if m.i == nil {
		return complex(m.r.At(r, c), 0)
	}
	if m.r == nil {
		return complex(0, m.i.At(r, c))
	}
	return complex(m.r.At(r, c), m.i.At(r, c))
}

// T performs an implicit transpose.
func (m Complex) T() cmat128.Matrix {
	if m.i == nil {
		return Complex{r: m.r.T()}
	}
	if m.r == nil {
		return Complex{i: m.i.T()}
	}
	return Complex{r: m.r.T(), i: m.i.T()}
}

// Real returns the real part of the receiver.
func (m Complex) Real() mat64.Matrix { return m.r }

// Imag returns the imaginary part of the receiver.
func (m Complex) Imag() mat64.Matrix { return m.i }

// Realer is a complex matrix that can return its real part.
type Realer interface {
	Real() mat64.Matrix
}

// Imager is a complex matrix that can return its imaginary part.
type Imager interface {
	Imag() mat64.Matrix
}

// Real is the real part of a complex matrix.
type Real struct{ Matrix cmat128.Matrix }

// NewReal returns a mat64.Matrix representing the real part of m. If m is a Realer,
// the real part is returned.
func NewReal(m cmat128.Matrix) mat64.Matrix {
	if m, ok := m.(Realer); ok {
		return m.Real()
	}
	return Real{m}
}

// Dims returns the number of rows and columns in the matrix.
func (m Real) Dims() (r, c int) { return m.Matrix.Dims() }

// At returns the element at row r, column c.
func (m Real) At(r, c int) float64 { return real(m.Matrix.At(r, c)) }

// T performs an implicit transpose.
func (m Real) T() mat64.Matrix { return Real{m.Matrix.T()} }

// Imag is the imaginary part of a complex matrix.
type Imag struct{ Matrix cmat128.Matrix }

// NewImage returns a mat64.Matrix representing the imaginary part of m. If m is an Imager,
// the imaginary part is returned.
func NewImag(m cmat128.Matrix) mat64.Matrix {
	if m, ok := m.(Imager); ok {
		return m.Imag()
	}
	return Imag{m}
}

// Dims returns the number of rows and columns in the matrix.
func (m Imag) Dims() (r, c int) { return m.Matrix.Dims() }

// At returns the element at row r, column c.
func (m Imag) At(r, c int) float64 { return imag(m.Matrix.At(r, c)) }

// T performs an implicit transpose.
func (m Imag) T() mat64.Matrix { return Imag{m.Matrix.T()} }
<commit_msg>conv: rename (r,c) params to (i,j)<commit_after>// Copyright ©2015 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package conv provides matrix type interconversion utilities.
package conv

import (
	"github.com/gonum/matrix"
	"github.com/gonum/matrix/cmat128"
	"github.com/gonum/matrix/mat64"
)

// Complex is a complex matrix constructed from two real matrices.
type Complex struct {
	// r and i are not exposed to ensure that
	// their dimensions can not be altered by
	// clients behind our back.
	r, i mat64.Matrix
}

var (
	_ Realer = Complex{}
	_ Imager = Complex{}
)

// NewComplex returns a complex matrix constructed from r and i. At least one of
// r or i must be non-nil otherwise NewComples will panic. If one of the inputs
// is nil, that part of the complex number will be zero when returned by At.
// If both are non-nil but differ in their sizes, NewComplex will panic.
func NewComplex(r, i mat64.Matrix) Complex {
	if r == nil && i == nil {
		panic("conv: no matrix")
	} else if r != nil && i != nil {
		rr, rc := r.Dims()
		ir, ic := i.Dims()
		if rr != ir || rc != ic {
			panic(matrix.ErrShape)
		}
	}
	return Complex{r: r, i: i}
}

// Dims returns the number of rows and columns in the matrix.
func (m Complex) Dims() (r, c int) {
	if m.r == nil {
		return m.i.Dims()
	}
	return m.r.Dims()
}

// At returns the element at row i, column j.
func (m Complex) At(i, j int) complex128 {
	if m.i == nil {
		return complex(m.r.At(i, j), 0)
	}
	if m.r == nil {
		return complex(0, m.i.At(i, j))
	}
	return complex(m.r.At(i, j), m.i.At(i, j))
}

// T performs an implicit transpose.
func (m Complex) T() cmat128.Matrix {
	if m.i == nil {
		return Complex{r: m.r.T()}
	}
	if m.r == nil {
		return Complex{i: m.i.T()}
	}
	return Complex{r: m.r.T(), i: m.i.T()}
}

// Real returns the real part of the receiver.
func (m Complex) Real() mat64.Matrix { return m.r }

// Imag returns the imaginary part of the receiver.
func (m Complex) Imag() mat64.Matrix { return m.i }

// Realer is a complex matrix that can return its real part.
type Realer interface {
	Real() mat64.Matrix
}

// Imager is a complex matrix that can return its imaginary part.
type Imager interface {
	Imag() mat64.Matrix
}

// Real is the real part of a complex matrix.
type Real struct{ Matrix cmat128.Matrix }

// NewReal returns a mat64.Matrix representing the real part of m. If m is a Realer,
// the real part is returned.
func NewReal(m cmat128.Matrix) mat64.Matrix {
	if m, ok := m.(Realer); ok {
		return m.Real()
	}
	return Real{m}
}

// Dims returns the number of rows and columns in the matrix.
func (m Real) Dims() (r, c int) { return m.Matrix.Dims() }

// At returns the element at row i, column j.
func (m Real) At(i, j int) float64 { return real(m.Matrix.At(i, j)) }

// T performs an implicit transpose.
func (m Real) T() mat64.Matrix { return Real{m.Matrix.T()} }

// Imag is the imaginary part of a complex matrix.
type Imag struct{ Matrix cmat128.Matrix }

// NewImage returns a mat64.Matrix representing the imaginary part of m. If m is an Imager,
// the imaginary part is returned.
func NewImag(m cmat128.Matrix) mat64.Matrix {
	if m, ok := m.(Imager); ok {
		return m.Imag()
	}
	return Imag{m}
}

// Dims returns the number of rows and columns in the matrix.
func (m Imag) Dims() (r, c int) { return m.Matrix.Dims() }

// At returns the element at row i, column j.
func (m Imag) At(i, j int) float64 { return imag(m.Matrix.At(i, j)) }

// T performs an implicit transpose.
func (m Imag) T() mat64.Matrix { return Imag{m.Matrix.T()} }
<|endoftext|>