<commit_before>package kontrolclient

import (
	"container/list"
	"errors"
	"net/url"
	"sync"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite"
	"github.com/koding/kite/protocol"
)

// Returned from GetKites when query matches no kites.
var ErrNoKitesAvailable = errors.New("no kites availabile")

// Kontrol is a client for registering and querying Kontrol Kite.
type Kontrol struct {
	*kite.RemoteKite

	LocalKite *kite.Kite

	// used for synchronizing methods that needs to be called after
	// successful connection.
	ready chan bool

	// Watchers are saved here to re-watch on reconnect.
	watchers      *list.List
	watchersMutex sync.RWMutex
}

// NewKontrol returns a pointer to new Kontrol instance.
func New(k *kite.Kite) *Kontrol {
	if k.Config.KontrolURL == nil {
		panic("no kontrol URL given in config")
	}

	remoteKite := k.NewRemoteKite(k.Config.KontrolURL)
	remoteKite.Kite = protocol.Kite{Name: "kontrol"} // for logging purposes
	remoteKite.Authentication = &kite.Authentication{
		Type: "kiteKey",
		Key:  k.Config.KiteKey,
	}

	kontrol := &Kontrol{
		RemoteKite: remoteKite,
		LocalKite:  k,
		ready:      make(chan bool),
		watchers:   list.New(),
	}

	var once sync.Once

	kontrol.OnConnect(func() {
		k.Log.Info("Connected to Kontrol ")

		// signal all other methods that are listening on this channel, that we
		// are ready.
		once.Do(func() { close(kontrol.ready) })

		// Re-register existing watchers.
		kontrol.watchersMutex.RLock()
		for e := kontrol.watchers.Front(); e != nil; e = e.Next() {
			watcher := e.Value.(*Watcher)
			if err := watcher.rewatch(); err != nil {
				kontrol.Log.Error("Cannot rewatch query: %+v", watcher)
			}
		}
		kontrol.watchersMutex.RUnlock()
	})

	return kontrol
}

type registerResult struct {
	URL *url.URL
}

// Register registers current Kite to Kontrol. After registration other Kites
// can find it via GetKites() method.
func (k *Kontrol) Register(kiteURL *url.URL) (*registerResult, error) {
	<-k.ready

	args := protocol.RegsiterArgs{
		URL: kiteURL.String(),
	}

	response, err := k.RemoteKite.Tell("register", args)
	if err != nil {
		return nil, err
	}

	var rr protocol.RegisterResult
	err = response.Unmarshal(&rr)
	if err != nil {
		return nil, err
	}

	k.Log.Info("Registered to kontrol with URL: %s", rr.URL)

	parsed, err := url.Parse(rr.URL)
	if err != nil {
		k.Log.Error("Cannot parse registered URL: %s", err.Error())
	}

	return &registerResult{parsed}, nil
}

// WatchKites watches for Kites that matches the query. The onEvent functions
// is called for current kites and every new kite event.
func (k *Kontrol) WatchKites(query protocol.KontrolQuery, onEvent EventHandler) (*Watcher, error) {
	<-k.ready

	watcherID, err := k.watchKites(query, onEvent)
	if err != nil {
		return nil, err
	}

	return k.newWatcher(watcherID, &query, onEvent), nil
}

func (k *Kontrol) eventCallbackHandler(onEvent EventHandler) kite.Callback {
	return func(r *kite.Request) {
		var returnEvent *Event
		var returnError error

		args := r.Args.MustSliceOfLength(2)

		// Unmarshal event argument
		if args[0] != nil {
			var event = Event{localKite: k.LocalKite}
			err := args[0].Unmarshal(&event)
			if err != nil {
				k.Log.Error(err.Error())
				return
			}
			returnEvent = &event
		}

		// Unmarshal error argument
		if args[1] != nil {
			var kiteErr kite.Error
			err := args[1].Unmarshal(&kiteErr)
			if err != nil {
				k.Log.Error(err.Error())
				return
			}
			returnError = &kiteErr
		}

		onEvent(returnEvent, returnError)
	}
}

func (k *Kontrol) watchKites(query protocol.KontrolQuery, onEvent EventHandler) (watcherID string, err error) {
	remoteKites, watcherID, err := k.getKites(query, k.eventCallbackHandler(onEvent))
	if err != nil && err != ErrNoKitesAvailable {
		return "", err // return only when something really happened
	}

	// also put the current kites to the eventChan.
	for _, remoteKite := range remoteKites {
		event := Event{
			KiteEvent: protocol.KiteEvent{
				Action: protocol.Register,
				Kite:   remoteKite.Kite,
				Token:  remoteKite.Authentication.Key,
			},
			localKite: k.LocalKite,
		}

		onEvent(&event, nil)
	}

	return watcherID, nil
}

// GetKites returns the list of Kites matching the query. The returned list
// contains ready to connect RemoteKite instances. The caller must connect
// with RemoteKite.Dial() before using each Kite. An error is returned when no
// kites are available.
func (k *Kontrol) GetKites(query protocol.KontrolQuery) ([]*kite.RemoteKite, error) {
	remoteKites, _, err := k.getKites(query)
	if err != nil {
		return nil, err
	}

	if len(remoteKites) == 0 {
		return nil, ErrNoKitesAvailable
	}

	return remoteKites, nil
}

// used internally for GetKites() and WatchKites()
func (k *Kontrol) getKites(args ...interface{}) (kites []*kite.RemoteKite, watcherID string, err error) {
	<-k.ready

	response, err := k.RemoteKite.Tell("getKites", args...)
	if err != nil {
		return nil, "", err
	}

	var result = new(protocol.GetKitesResult)
	err = response.Unmarshal(&result)
	if err != nil {
		return nil, "", err
	}

	remoteKites := make([]*kite.RemoteKite, len(result.Kites))
	for i, currentKite := range result.Kites {
		_, err := jwt.Parse(currentKite.Token, k.LocalKite.RSAKey)
		if err != nil {
			return nil, result.WatcherID, err
		}

		// exp := time.Unix(int64(token.Claims["exp"].(float64)), 0).UTC()
		auth := &kite.Authentication{
			Type: "token",
			Key:  currentKite.Token,
			// validUntil: &exp,
		}

		parsed, err := url.Parse(currentKite.URL)
		if err != nil {
			k.Log.Error("invalid url came from kontrol", currentKite.URL)
		}

		remoteKites[i] = k.LocalKite.NewRemoteKiteString(currentKite.URL)
		remoteKites[i].Kite = currentKite.Kite
		remoteKites[i].URL = parsed
		remoteKites[i].Authentication = auth
	}

	return remoteKites, result.WatcherID, nil
}

// GetToken is used to get a new token for a single Kite.
func (k *Kontrol) GetToken(kite *protocol.Kite) (string, error) {
	<-k.ready

	result, err := k.RemoteKite.Tell("getToken", kite)
	if err != nil {
		return "", err
	}

	var tkn string
	err = result.Unmarshal(&tkn)
	if err != nil {
		return "", err
	}

	return tkn, nil
}
<commit_msg>update kontrol client comment<commit_after>package kontrolclient

import (
	"container/list"
	"errors"
	"net/url"
	"sync"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite"
	"github.com/koding/kite/protocol"
)

// Returned from GetKites when query matches no kites.
var ErrNoKitesAvailable = errors.New("no kites availabile")

// Kontrol is a client for registering and querying Kontrol Kite.
type Kontrol struct {
	*kite.RemoteKite

	LocalKite *kite.Kite

	// used for synchronizing methods that needs to be called after
	// successful connection.
	ready chan bool

	// Watchers are saved here to re-watch on reconnect.
	watchers      *list.List
	watchersMutex sync.RWMutex
}

// NewKontrol returns a pointer to new Kontrol instance.
func New(k *kite.Kite) *Kontrol {
	if k.Config.KontrolURL == nil {
		panic("no kontrol URL given in config")
	}

	remoteKite := k.NewRemoteKite(k.Config.KontrolURL)
	remoteKite.Kite = protocol.Kite{Name: "kontrol"} // for logging purposes
	remoteKite.Authentication = &kite.Authentication{
		Type: "kiteKey",
		Key:  k.Config.KiteKey,
	}

	kontrol := &Kontrol{
		RemoteKite: remoteKite,
		LocalKite:  k,
		ready:      make(chan bool),
		watchers:   list.New(),
	}

	var once sync.Once

	kontrol.OnConnect(func() {
		k.Log.Info("Connected to Kontrol ")

		// signal all other methods that are listening on this channel, that we
		// are ready.
		once.Do(func() { close(kontrol.ready) })

		// Re-register existing watchers.
		kontrol.watchersMutex.RLock()
		for e := kontrol.watchers.Front(); e != nil; e = e.Next() {
			watcher := e.Value.(*Watcher)
			if err := watcher.rewatch(); err != nil {
				kontrol.Log.Error("Cannot rewatch query: %+v", watcher)
			}
		}
		kontrol.watchersMutex.RUnlock()
	})

	return kontrol
}

type registerResult struct {
	URL *url.URL
}

// Register registers current Kite to Kontrol. After registration other Kites
// can find it via GetKites() method.
//
// This method does not handle the reconnection case. If you want to keep
// registered to kontrol, use kite/registration package.
func (k *Kontrol) Register(kiteURL *url.URL) (*registerResult, error) {
	<-k.ready

	args := protocol.RegsiterArgs{
		URL: kiteURL.String(),
	}

	response, err := k.RemoteKite.Tell("register", args)
	if err != nil {
		return nil, err
	}

	var rr protocol.RegisterResult
	err = response.Unmarshal(&rr)
	if err != nil {
		return nil, err
	}

	k.Log.Info("Registered to kontrol with URL: %s", rr.URL)

	parsed, err := url.Parse(rr.URL)
	if err != nil {
		k.Log.Error("Cannot parse registered URL: %s", err.Error())
	}

	return &registerResult{parsed}, nil
}

// WatchKites watches for Kites that matches the query. The onEvent functions
// is called for current kites and every new kite event.
func (k *Kontrol) WatchKites(query protocol.KontrolQuery, onEvent EventHandler) (*Watcher, error) {
	<-k.ready

	watcherID, err := k.watchKites(query, onEvent)
	if err != nil {
		return nil, err
	}

	return k.newWatcher(watcherID, &query, onEvent), nil
}

func (k *Kontrol) eventCallbackHandler(onEvent EventHandler) kite.Callback {
	return func(r *kite.Request) {
		var returnEvent *Event
		var returnError error

		args := r.Args.MustSliceOfLength(2)

		// Unmarshal event argument
		if args[0] != nil {
			var event = Event{localKite: k.LocalKite}
			err := args[0].Unmarshal(&event)
			if err != nil {
				k.Log.Error(err.Error())
				return
			}
			returnEvent = &event
		}

		// Unmarshal error argument
		if args[1] != nil {
			var kiteErr kite.Error
			err := args[1].Unmarshal(&kiteErr)
			if err != nil {
				k.Log.Error(err.Error())
				return
			}
			returnError = &kiteErr
		}

		onEvent(returnEvent, returnError)
	}
}

func (k *Kontrol) watchKites(query protocol.KontrolQuery, onEvent EventHandler) (watcherID string, err error) {
	remoteKites, watcherID, err := k.getKites(query, k.eventCallbackHandler(onEvent))
	if err != nil && err != ErrNoKitesAvailable {
		return "", err // return only when something really happened
	}

	// also put the current kites to the eventChan.
	for _, remoteKite := range remoteKites {
		event := Event{
			KiteEvent: protocol.KiteEvent{
				Action: protocol.Register,
				Kite:   remoteKite.Kite,
				Token:  remoteKite.Authentication.Key,
			},
			localKite: k.LocalKite,
		}

		onEvent(&event, nil)
	}

	return watcherID, nil
}

// GetKites returns the list of Kites matching the query. The returned list
// contains ready to connect RemoteKite instances. The caller must connect
// with RemoteKite.Dial() before using each Kite. An error is returned when no
// kites are available.
func (k *Kontrol) GetKites(query protocol.KontrolQuery) ([]*kite.RemoteKite, error) {
	remoteKites, _, err := k.getKites(query)
	if err != nil {
		return nil, err
	}

	if len(remoteKites) == 0 {
		return nil, ErrNoKitesAvailable
	}

	return remoteKites, nil
}

// used internally for GetKites() and WatchKites()
func (k *Kontrol) getKites(args ...interface{}) (kites []*kite.RemoteKite, watcherID string, err error) {
	<-k.ready

	response, err := k.RemoteKite.Tell("getKites", args...)
	if err != nil {
		return nil, "", err
	}

	var result = new(protocol.GetKitesResult)
	err = response.Unmarshal(&result)
	if err != nil {
		return nil, "", err
	}

	remoteKites := make([]*kite.RemoteKite, len(result.Kites))
	for i, currentKite := range result.Kites {
		_, err := jwt.Parse(currentKite.Token, k.LocalKite.RSAKey)
		if err != nil {
			return nil, result.WatcherID, err
		}

		// exp := time.Unix(int64(token.Claims["exp"].(float64)), 0).UTC()
		auth := &kite.Authentication{
			Type: "token",
			Key:  currentKite.Token,
			// validUntil: &exp,
		}

		parsed, err := url.Parse(currentKite.URL)
		if err != nil {
			k.Log.Error("invalid url came from kontrol", currentKite.URL)
		}

		remoteKites[i] = k.LocalKite.NewRemoteKiteString(currentKite.URL)
		remoteKites[i].Kite = currentKite.Kite
		remoteKites[i].URL = parsed
		remoteKites[i].Authentication = auth
	}

	return remoteKites, result.WatcherID, nil
}

// GetToken is used to get a new token for a single Kite.
func (k *Kontrol) GetToken(kite *protocol.Kite) (string, error) {
	<-k.ready

	result, err := k.RemoteKite.Tell("getToken", kite)
	if err != nil {
		return "", err
	}

	var tkn string
	err = result.Unmarshal(&tkn)
	if err != nil {
		return "", err
	}

	return tkn, nil
}
<|endoftext|>


<commit_before>package dbus

import (
	"encoding/binary"
	"io"
	"io/ioutil"
	"testing"
	"time"
)

func TestSessionBus(t *testing.T) {
	_, err := SessionBus()
	if err != nil {
		t.Error(err)
	}
}

func TestSystemBus(t *testing.T) {
	_, err := SystemBus()
	if err != nil {
		t.Error(err)
	}
}

func ExampleSystemBusPrivate() {
	setupPrivateSystemBus := func() (conn *Conn, err error) {
		conn, err = SystemBusPrivate()
		if err != nil {
			return nil, err
		}
		if err = conn.Auth(nil); err != nil {
			conn.Close()
			conn = nil
			return
		}
		if err = conn.Hello(); err != nil {
			conn.Close()
			conn = nil
		}
		return conn, nil // success
	}
	_, _ = setupPrivateSystemBus()
}

func TestSend(t *testing.T) {
	bus, err := SessionBus()
	if err != nil {
		t.Fatal(err)
	}
	ch := make(chan *Call, 1)
	msg := &Message{
		Type:  TypeMethodCall,
		Flags: 0,
		Headers: map[HeaderField]Variant{
			FieldDestination: MakeVariant(bus.Names()[0]),
			FieldPath:        MakeVariant(ObjectPath("/org/freedesktop/DBus")),
			FieldInterface:   MakeVariant("org.freedesktop.DBus.Peer"),
			FieldMember:      MakeVariant("Ping"),
		},
	}
	call := bus.Send(msg, ch)
	<-ch
	if call.Err != nil {
		t.Error(call.Err)
	}
}

func TestFlagNoReplyExpectedSend(t *testing.T) {
	bus, err := SessionBus()
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan struct{})
	go func() {
		bus.BusObject().Call("org.freedesktop.DBus.ListNames", FlagNoReplyExpected)
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(1 * time.Second):
		t.Error("Failed to announce that the call was done")
	}
}

func TestRemoveSignal(t *testing.T) {
	bus, err := NewConn(nil)
	if err != nil {
		t.Error(err)
	}
	signals := bus.signalHandler.(*defaultSignalHandler).signals
	ch := make(chan *Signal)
	ch2 := make(chan *Signal)
	for _, ch := range []chan *Signal{ch, ch2, ch, ch2, ch2, ch} {
		bus.Signal(ch)
	}
	signals = bus.signalHandler.(*defaultSignalHandler).signals
	if len(signals) != 6 {
		t.Errorf("remove signal: signals length not equal: got '%d', want '6'", len(signals))
	}
	bus.RemoveSignal(ch)
	signals = bus.signalHandler.(*defaultSignalHandler).signals
	if len(signals) != 3 {
		t.Errorf("remove signal: signals length not equal: got '%d', want '3'", len(signals))
	}
	signals = bus.signalHandler.(*defaultSignalHandler).signals
	for _, bch := range signals {
		if bch != ch2 {
			t.Errorf("remove signal: removed signal present: got '%v', want '%v'", bch, ch2)
		}
	}
}

type rwc struct {
	io.Reader
	io.Writer
}

func (rwc) Close() error { return nil }

type fakeAuth struct {
}

func (fakeAuth) FirstData() (name, resp []byte, status AuthStatus) {
	return []byte("name"), []byte("resp"), AuthOk
}

func (fakeAuth) HandleData(data []byte) (resp []byte, status AuthStatus) {
	return nil, AuthOk
}

func TestCloseBeforeSignal(t *testing.T) {
	reader, pipewriter := io.Pipe()
	defer pipewriter.Close()
	defer reader.Close()

	bus, err := NewConn(rwc{Reader: reader, Writer: ioutil.Discard})
	if err != nil {
		t.Fatal(err)
	}
	// give ch a buffer so sends won't block
	ch := make(chan *Signal, 1)
	bus.Signal(ch)

	go func() {
		_, err := pipewriter.Write([]byte("REJECTED name\r\nOK myuuid\r\n"))
		if err != nil {
			t.Errorf("error writing to pipe: %v", err)
		}
	}()

	err = bus.Auth([]Auth{fakeAuth{}})
	if err != nil {
		t.Fatal(err)
	}

	err = bus.Close()
	if err != nil {
		t.Fatal(err)
	}

	msg := &Message{
		Type: TypeSignal,
		Headers: map[HeaderField]Variant{
			FieldInterface: MakeVariant("foo.bar"),
			FieldMember:    MakeVariant("bar"),
			FieldPath:      MakeVariant(ObjectPath("/baz")),
		},
	}
	err = msg.EncodeTo(pipewriter, binary.LittleEndian)
	if err != nil {
		t.Fatal(err)
	}
}

func TestAddAndRemoveMatchSignal(t *testing.T) {
	conn, err := SessionBusPrivate()
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()

	if err = conn.Auth(nil); err != nil {
		t.Fatal(err)
	}
	if err = conn.Hello(); err != nil {
		t.Fatal(err)
	}

	sigc := make(chan *Signal, 1)
	conn.Signal(sigc)

	// subscribe to a made up signal name and emit one of the type
	if err = conn.AddMatchSignal(
		WithMatchInterface("org.test"),
		WithMatchMember("Test"),
	); err != nil {
		t.Fatal(err)
	}
	if err = conn.Emit("/", "org.test.Test"); err != nil {
		t.Fatal(err)
	}
	if sig := waitSignal(sigc, "org.test.Test", time.Second); sig == nil {
		t.Fatal("signal receive timed out")
	}

	// unsubscribe from the signal and check that is not delivered anymore
	if err = conn.RemoveMatchSignal(
		WithMatchInterface("org.test"),
		WithMatchMember("Test"),
	); err != nil {
		t.Fatal(err)
	}
	if err = conn.Emit("/", "org.test.Test"); err != nil {
		t.Fatal(err)
	}
	if sig := waitSignal(sigc, "org.test.Test", time.Second); sig != nil {
		t.Fatalf("unsubscribed from %q signal, but received %#v", "org.test.Test", sig)
	}
}

func waitSignal(sigc <-chan *Signal, name string, timeout time.Duration) *Signal {
	for {
		select {
		case sig := <-sigc:
			if sig.Name == name {
				return sig
			}
		case <-time.After(timeout):
			return nil
		}
	}
}

type server struct{}

func (server) Double(i int64) (int64, *Error) {
	return 2 * i, nil
}

func BenchmarkCall(b *testing.B) {
	b.StopTimer()
	b.ReportAllocs()
	var s string
	bus, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}
	name := bus.Names()[0]
	obj := bus.BusObject()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		err := obj.Call("org.freedesktop.DBus.GetNameOwner", 0, name).Store(&s)
		if err != nil {
			b.Fatal(err)
		}
		if s != name {
			b.Errorf("got %s, wanted %s", s, name)
		}
	}
}

func BenchmarkCallAsync(b *testing.B) {
	b.StopTimer()
	b.ReportAllocs()
	bus, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}
	name := bus.Names()[0]
	obj := bus.BusObject()
	c := make(chan *Call, 50)
	done := make(chan struct{})
	go func() {
		for i := 0; i < b.N; i++ {
			v := <-c
			if v.Err != nil {
				b.Error(v.Err)
			}
			s := v.Body[0].(string)
			if s != name {
				b.Errorf("got %s, wanted %s", s, name)
			}
		}
		close(done)
	}()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		obj.Go("org.freedesktop.DBus.GetNameOwner", 0, c, name)
	}
	<-done
}

func BenchmarkServe(b *testing.B) {
	b.StopTimer()
	srv, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}
	cli, err := SessionBusPrivate()
	if err != nil {
		b.Fatal(err)
	}
	if err = cli.Auth(nil); err != nil {
		b.Fatal(err)
	}
	if err = cli.Hello(); err != nil {
		b.Fatal(err)
	}
	benchmarkServe(b, srv, cli)
}

func BenchmarkServeAsync(b *testing.B) {
	b.StopTimer()
	srv, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}
	cli, err := SessionBusPrivate()
	if err != nil {
		b.Fatal(err)
	}
	if err = cli.Auth(nil); err != nil {
		b.Fatal(err)
	}
	if err = cli.Hello(); err != nil {
		b.Fatal(err)
	}
	benchmarkServeAsync(b, srv, cli)
}

func BenchmarkServeSameConn(b *testing.B) {
	b.StopTimer()
	bus, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}

	benchmarkServe(b, bus, bus)
}

func BenchmarkServeSameConnAsync(b *testing.B) {
	b.StopTimer()
	bus, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}

	benchmarkServeAsync(b, bus, bus)
}

func benchmarkServe(b *testing.B, srv, cli *Conn) {
	var r int64
	var err error
	dest := srv.Names()[0]
	srv.Export(server{}, "/org/guelfey/DBus/Test", "org.guelfey.DBus.Test")
	obj := cli.Object(dest, "/org/guelfey/DBus/Test")
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		err = obj.Call("org.guelfey.DBus.Test.Double", 0, int64(i)).Store(&r)
		if err != nil {
			b.Fatal(err)
		}
		if r != 2*int64(i) {
			b.Errorf("got %d, wanted %d", r, 2*int64(i))
		}
	}
}

func benchmarkServeAsync(b *testing.B, srv, cli *Conn) {
	dest := srv.Names()[0]
	srv.Export(server{}, "/org/guelfey/DBus/Test", "org.guelfey.DBus.Test")
	obj := cli.Object(dest, "/org/guelfey/DBus/Test")
	c := make(chan *Call, 50)
	done := make(chan struct{})
	go func() {
		for i := 0; i < b.N; i++ {
			v := <-c
			if v.Err != nil {
				b.Fatal(v.Err)
			}
			i, r := v.Args[0].(int64), v.Body[0].(int64)
			if 2*i != r {
				b.Errorf("got %d, wanted %d", r, 2*i)
			}
		}
		close(done)
	}()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		obj.Go("org.guelfey.DBus.Test.Double", 0, c, int64(i))
	}
	<-done
}

func TestGetKey(t *testing.T) {
	keys := "host=1.2.3.4,port=5678,family=ipv4"
	if host := getKey(keys, "host"); host != "1.2.3.4" {
		t.Error(`Expected "1.2.3.4", got`, host)
	}
	if port := getKey(keys, "port"); port != "5678" {
		t.Error(`Expected "5678", got`, port)
	}
	if family := getKey(keys, "family"); family != "ipv4" {
		t.Error(`Expected "ipv4", got`, family)
	}
}
<commit_msg>tests: add a test for closing a signal channel after calling RemoveSignal<commit_after>package dbus

import (
	"encoding/binary"
	"io"
	"io/ioutil"
	"testing"
	"time"
)

func TestSessionBus(t *testing.T) {
	_, err := SessionBus()
	if err != nil {
		t.Error(err)
	}
}

func TestSystemBus(t *testing.T) {
	_, err := SystemBus()
	if err != nil {
		t.Error(err)
	}
}

func ExampleSystemBusPrivate() {
	setupPrivateSystemBus := func() (conn *Conn, err error) {
		conn, err = SystemBusPrivate()
		if err != nil {
			return nil, err
		}
		if err = conn.Auth(nil); err != nil {
			conn.Close()
			conn = nil
			return
		}
		if err = conn.Hello(); err != nil {
			conn.Close()
			conn = nil
		}
		return conn, nil // success
	}
	_, _ = setupPrivateSystemBus()
}

func TestSend(t *testing.T) {
	bus, err := SessionBus()
	if err != nil {
		t.Fatal(err)
	}
	ch := make(chan *Call, 1)
	msg := &Message{
		Type:  TypeMethodCall,
		Flags: 0,
		Headers: map[HeaderField]Variant{
			FieldDestination: MakeVariant(bus.Names()[0]),
			FieldPath:        MakeVariant(ObjectPath("/org/freedesktop/DBus")),
			FieldInterface:   MakeVariant("org.freedesktop.DBus.Peer"),
			FieldMember:      MakeVariant("Ping"),
		},
	}
	call := bus.Send(msg, ch)
	<-ch
	if call.Err != nil {
		t.Error(call.Err)
	}
}

func TestFlagNoReplyExpectedSend(t *testing.T) {
	bus, err := SessionBus()
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan struct{})
	go func() {
		bus.BusObject().Call("org.freedesktop.DBus.ListNames", FlagNoReplyExpected)
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(1 * time.Second):
		t.Error("Failed to announce that the call was done")
	}
}

func TestRemoveSignal(t *testing.T) {
	bus, err := NewConn(nil)
	if err != nil {
		t.Error(err)
	}
	signals := bus.signalHandler.(*defaultSignalHandler).signals
	ch := make(chan *Signal)
	ch2 := make(chan *Signal)
	for _, ch := range []chan *Signal{ch, ch2, ch, ch2, ch2, ch} {
		bus.Signal(ch)
	}
	signals = bus.signalHandler.(*defaultSignalHandler).signals
	if len(signals) != 6 {
		t.Errorf("remove signal: signals length not equal: got '%d', want '6'", len(signals))
	}
	bus.RemoveSignal(ch)
	signals = bus.signalHandler.(*defaultSignalHandler).signals
	if len(signals) != 3 {
		t.Errorf("remove signal: signals length not equal: got '%d', want '3'", len(signals))
	}
	signals = bus.signalHandler.(*defaultSignalHandler).signals
	for _, bch := range signals {
		if bch != ch2 {
			t.Errorf("remove signal: removed signal present: got '%v', want '%v'", bch, ch2)
		}
	}
}

type rwc struct {
	io.Reader
	io.Writer
}

func (rwc) Close() error { return nil }

type fakeAuth struct {
}

func (fakeAuth) FirstData() (name, resp []byte, status AuthStatus) {
	return []byte("name"), []byte("resp"), AuthOk
}

func (fakeAuth) HandleData(data []byte) (resp []byte, status AuthStatus) {
	return nil, AuthOk
}

func TestCloseBeforeSignal(t *testing.T) {
	reader, pipewriter := io.Pipe()
	defer pipewriter.Close()
	defer reader.Close()

	bus, err := NewConn(rwc{Reader: reader, Writer: ioutil.Discard})
	if err != nil {
		t.Fatal(err)
	}
	// give ch a buffer so sends won't block
	ch := make(chan *Signal, 1)
	bus.Signal(ch)

	go func() {
		_, err := pipewriter.Write([]byte("REJECTED name\r\nOK myuuid\r\n"))
		if err != nil {
			t.Errorf("error writing to pipe: %v", err)
		}
	}()

	err = bus.Auth([]Auth{fakeAuth{}})
	if err != nil {
		t.Fatal(err)
	}

	err = bus.Close()
	if err != nil {
		t.Fatal(err)
	}

	msg := &Message{
		Type: TypeSignal,
		Headers: map[HeaderField]Variant{
			FieldInterface: MakeVariant("foo.bar"),
			FieldMember:    MakeVariant("bar"),
			FieldPath:      MakeVariant(ObjectPath("/baz")),
		},
	}
	err = msg.EncodeTo(pipewriter, binary.LittleEndian)
	if err != nil {
		t.Fatal(err)
	}
}

func TestCloseChannelAfterRemoveSignal(t *testing.T) {
	bus, err := NewConn(nil)
	if err != nil {
		t.Fatal(err)
	}

	// Add an unbuffered signal channel
	ch := make(chan *Signal)
	bus.Signal(ch)

	// Send a signal
	msg := &Message{
		Type: TypeSignal,
		Headers: map[HeaderField]Variant{
			FieldInterface: MakeVariant("foo.bar"),
			FieldMember:    MakeVariant("bar"),
			FieldPath:      MakeVariant(ObjectPath("/baz")),
		},
	}
	bus.handleSignal(msg)

	// Remove and close the signal channel
	bus.RemoveSignal(ch)
	close(ch)
}

func TestAddAndRemoveMatchSignal(t *testing.T) {
	conn, err := SessionBusPrivate()
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()

	if err = conn.Auth(nil); err != nil {
		t.Fatal(err)
	}
	if err = conn.Hello(); err != nil {
		t.Fatal(err)
	}

	sigc := make(chan *Signal, 1)
	conn.Signal(sigc)

	// subscribe to a made up signal name and emit one of the type
	if err = conn.AddMatchSignal(
		WithMatchInterface("org.test"),
		WithMatchMember("Test"),
	); err != nil {
		t.Fatal(err)
	}
	if err = conn.Emit("/", "org.test.Test"); err != nil {
		t.Fatal(err)
	}
	if sig := waitSignal(sigc, "org.test.Test", time.Second); sig == nil {
		t.Fatal("signal receive timed out")
	}

	// unsubscribe from the signal and check that is not delivered anymore
	if err = conn.RemoveMatchSignal(
		WithMatchInterface("org.test"),
		WithMatchMember("Test"),
	); err != nil {
		t.Fatal(err)
	}
	if err = conn.Emit("/", "org.test.Test"); err != nil {
		t.Fatal(err)
	}
	if sig := waitSignal(sigc, "org.test.Test", time.Second); sig != nil {
		t.Fatalf("unsubscribed from %q signal, but received %#v", "org.test.Test", sig)
	}
}

func waitSignal(sigc <-chan *Signal, name string, timeout time.Duration) *Signal {
	for {
		select {
		case sig := <-sigc:
			if sig.Name == name {
				return sig
			}
		case <-time.After(timeout):
			return nil
		}
	}
}

type server struct{}

func (server) Double(i int64) (int64, *Error) {
	return 2 * i, nil
}

func BenchmarkCall(b *testing.B) {
	b.StopTimer()
	b.ReportAllocs()
	var s string
	bus, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}
	name := bus.Names()[0]
	obj := bus.BusObject()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		err := obj.Call("org.freedesktop.DBus.GetNameOwner", 0, name).Store(&s)
		if err != nil {
			b.Fatal(err)
		}
		if s != name {
			b.Errorf("got %s, wanted %s", s, name)
		}
	}
}

func BenchmarkCallAsync(b *testing.B) {
	b.StopTimer()
	b.ReportAllocs()
	bus, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}
	name := bus.Names()[0]
	obj := bus.BusObject()
	c := make(chan *Call, 50)
	done := make(chan struct{})
	go func() {
		for i := 0; i < b.N; i++ {
			v := <-c
			if v.Err != nil {
				b.Error(v.Err)
			}
			s := v.Body[0].(string)
			if s != name {
				b.Errorf("got %s, wanted %s", s, name)
			}
		}
		close(done)
	}()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		obj.Go("org.freedesktop.DBus.GetNameOwner", 0, c, name)
	}
	<-done
}

func BenchmarkServe(b *testing.B) {
	b.StopTimer()
	srv, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}
	cli, err := SessionBusPrivate()
	if err != nil {
		b.Fatal(err)
	}
	if err = cli.Auth(nil); err != nil {
		b.Fatal(err)
	}
	if err = cli.Hello(); err != nil {
		b.Fatal(err)
	}
	benchmarkServe(b, srv, cli)
}

func BenchmarkServeAsync(b *testing.B) {
	b.StopTimer()
	srv, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}
	cli, err := SessionBusPrivate()
	if err != nil {
		b.Fatal(err)
	}
	if err = cli.Auth(nil); err != nil {
		b.Fatal(err)
	}
	if err = cli.Hello(); err != nil {
		b.Fatal(err)
	}
	benchmarkServeAsync(b, srv, cli)
}

func BenchmarkServeSameConn(b *testing.B) {
	b.StopTimer()
	bus, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}

	benchmarkServe(b, bus, bus)
}

func BenchmarkServeSameConnAsync(b *testing.B) {
	b.StopTimer()
	bus, err := SessionBus()
	if err != nil {
		b.Fatal(err)
	}

	benchmarkServeAsync(b, bus, bus)
}

func benchmarkServe(b *testing.B, srv, cli *Conn) {
	var r int64
	var err error
	dest := srv.Names()[0]
	srv.Export(server{}, "/org/guelfey/DBus/Test", "org.guelfey.DBus.Test")
	obj := cli.Object(dest, "/org/guelfey/DBus/Test")
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		err = obj.Call("org.guelfey.DBus.Test.Double", 0, int64(i)).Store(&r)
		if err != nil {
			b.Fatal(err)
		}
		if r != 2*int64(i) {
			b.Errorf("got %d, wanted %d", r, 2*int64(i))
		}
	}
}

func benchmarkServeAsync(b *testing.B, srv, cli *Conn) {
	dest := srv.Names()[0]
	srv.Export(server{}, "/org/guelfey/DBus/Test", "org.guelfey.DBus.Test")
	obj := cli.Object(dest, "/org/guelfey/DBus/Test")
	c := make(chan *Call, 50)
	done := make(chan struct{})
	go func() {
		for i := 0; i < b.N; i++ {
			v := <-c
			if v.Err != nil {
				b.Fatal(v.Err)
			}
			i, r := v.Args[0].(int64), v.Body[0].(int64)
			if 2*i != r {
				b.Errorf("got %d, wanted %d", r, 2*i)
			}
		}
		close(done)
	}()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		obj.Go("org.guelfey.DBus.Test.Double", 0, c, int64(i))
	}
	<-done
}

func TestGetKey(t *testing.T) {
	keys := "host=1.2.3.4,port=5678,family=ipv4"
	if host := getKey(keys, "host"); host != "1.2.3.4" {
		t.Error(`Expected "1.2.3.4", got`, host)
	}
	if port := getKey(keys, "port"); port != "5678" {
		t.Error(`Expected "5678", got`, port)
	}
	if family := getKey(keys, "family"); family != "ipv4" {
		t.Error(`Expected "ipv4", got`, family)
	}
}
<|endoftext|>


<commit_before>package floatutils_test
<commit_msg>Removed empty file<commit_after><|endoftext|>


<commit_before>package mit

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var parensRegex = regexp.MustCompile("\\((.*?)\\)")

// A ScheduledMeeting struct contains details about a scheduled meeting in a TimeInfo. Each TimeInfo can have multiple ScheduledMeetings.
type ScheduledMeeting struct {
	Weekdays     []time.Weekday
	StartSeconds int
	EndSeconds   int
}

// A TimeInfo struct contains details about when a class's section meets.
type TimeInfo struct {
	Meetings []ScheduledMeeting
	BeginsOn time.Time
	EndsOn   time.Time
}

func parseTime(timeString string, forceAM bool) (int, error) {
	// sometimes, for whatever reason, use a colon instead of the normal dot as a separator
	// this seems to just be some advising seminars
	// replace these colons with a dot so that we can handle them
	timeString = strings.Replace(timeString, ":", ".", -1)

	// some classes also like to just randomly add a "pm"
	// looking at you WGS.228 and 15.389
	isPM := false
	if strings.Contains(timeString, "pm") || strings.Contains(timeString, "PM") {
		timeString = strings.Replace(strings.Replace(timeString, "pm", "", -1), "PM", "", -1)
		isPM = true
	}

	// also check for just "p"
	if strings.Contains(timeString, "p") || strings.Contains(timeString, "P") {
		timeString = strings.Replace(strings.Replace(timeString, "p", "", -1), "P", "", -1)
		isPM = true
	}

	// it's a time like "4" or "5.30"
	parts := strings.Split(timeString, ".")
	if len(parts) > 2 {
		return 0, fmt.Errorf("mit: parseTime: time string '%s' had too many dots", timeString)
	}

	hour, err := strconv.Atoi(strings.TrimSpace(parts[0]))
	if err != nil {
		return 0, err
	}

	minute := 0

	if len(parts) == 2 {
		minute, err = strconv.Atoi(strings.TrimSpace(parts[1]))
		if err != nil {
			return 0, err
		}
	}

	if forceAM || (!isPM && hour >= 8 && hour <= 12) {
		// assume AM
	} else {
		// assume PM
		hour += 12
		if hour == 24 {
			// 12 PM is actually noon
			hour = 12
		}
	}

	timeSeconds := ((hour * 60) + minute) * 60

	return timeSeconds, nil
}

func assignYearForDate(date time.Time, termInfo TermInfo) time.Time {
	return date.AddDate(termInfo.FirstDayOfClasses.Year(), 0, 0)
}

// ParseScheduledMeeting parses the given time info string, in a format like "MW4-5.30"
func ParseScheduledMeeting(scheduledMeetingString string, termInfo TermInfo) (*ScheduledMeeting, *time.Time, *time.Time, error) {
	scheduledMeeting := ScheduledMeeting{
		Weekdays: []time.Weekday{},
	}

	var beginsOn *time.Time
	var endsOn *time.Time

	scheduledMeetingString = strings.ToUpper(strings.TrimSpace(strings.Replace(scheduledMeetingString, "*", "", -1)))

	if scheduledMeetingString == "SUBJECT CANCELLED" || scheduledMeetingString == "SECTION CANCELLED" || scheduledMeetingString == "TO BE ARRANGED" || scheduledMeetingString == "TBA" || scheduledMeetingString == "TBD" || scheduledMeetingString == "ARRANGED" || scheduledMeetingString == "SCHEDULED INDIVIDUALLY" || scheduledMeetingString == "SCHECULED INDIVIDUALLY" || scheduledMeetingString == "UNSCHEDULED" {
		// oof
		return nil, nil, nil, nil
	}

	// HACK: some classes use "TH" for thursday instead of a single character
	// for example, 12.803 has "TTH 10:30 - 12 PM"
	// we work around this by replacing "TH" with the correct code, "R"
	// this works because H is not used for anything else
	scheduledMeetingString = strings.ReplaceAll(scheduledMeetingString, "TH", "R")

	parsedATime := false

	if strings.Contains(scheduledMeetingString, "(") {
		// there's a thing in parentheses
		// for example "TR10.30-12 (BEGINS OCT 21)"

		// it could be an evening class
		// for example "TR EVE (4-6 PM)" or "W EVE (4-6.30 PM)"
		isEvening := (strings.Contains(scheduledMeetingString, "EVE"))

		matches := parensRegex.FindAllStringSubmatch(scheduledMeetingString, -1)
		for i, match := range matches {
			info := match[1]

			if i == 0 && isEvening {
				// it's an evening class
				// the first paren is a time
				info = strings.Replace(info, " PM", "", -1)

				subScheduledMeeting, _, _, err := ParseScheduledMeeting(info, termInfo)
				if err != nil {
					return nil, nil, nil, err
				}

				if subScheduledMeeting.StartSeconds < 12*60*60 {
					subScheduledMeeting.StartSeconds += 12 * 60 * 60
				} else if subScheduledMeeting.EndSeconds < 12*60*60 {
					subScheduledMeeting.EndSeconds += 12 * 60 * 60
				}

				scheduledMeeting.StartSeconds = subScheduledMeeting.StartSeconds
				scheduledMeeting.EndSeconds = subScheduledMeeting.EndSeconds

				scheduledMeetingString = strings.Replace(scheduledMeetingString, "EVE", "", -1)
				scheduledMeetingString = parensRegex.ReplaceAllString(scheduledMeetingString, "")

				parsedATime = true

				continue
			}

			if strings.HasPrefix(info, "BEGINS") {
				dateString := strings.Replace(info, "BEGINS ", "", -1)
				parsedDate, err := time.Parse("Jan 2", dateString)
				if err != nil {
					return nil, nil, nil, err
				}
				parsedDate = assignYearForDate(parsedDate, termInfo)
				beginsOn = &parsedDate
			} else if strings.HasPrefix(info, "END") {
				dateString := strings.Replace(info, "ENDS ", "", -1)
				parsedDate, err := time.Parse("Jan 2", dateString)
				if err != nil {
					return nil, nil, nil, err
				}
				parsedDate = assignYearForDate(parsedDate, termInfo)
				endsOn = &parsedDate
			} else if strings.HasPrefix(info, "MEETS") {
				// it's a message in the format "MEETS 4/7 TO 5/14"
				dateString := strings.Replace(info, "MEETS ", "", -1)
				parts := strings.Split(dateString, "TO")

				parsedBegin, err := time.Parse("1/2", strings.TrimSpace(parts[0]))
				if err != nil {
					return nil, nil, nil, err
				}
				parsedBegin = assignYearForDate(parsedBegin, termInfo)
				beginsOn = &parsedBegin

				parsedEnd, err := time.Parse("1/2", strings.TrimSpace(parts[1]))
				if err != nil {
					return nil, nil, nil, err
				}
				parsedEnd = assignYearForDate(parsedEnd, termInfo)
				endsOn = &parsedEnd
			} else if strings.HasPrefix(info, "LIMITED") {
				// it's about limited enrollment, which we don't care about
				// just ignore it
			} else {
				// some classes are not entered into the catalog correctly
				// they have a format like T(11-12:30)
				// if we can parse the time in parenthesis, assume that's the problem
				parsedInfo, _, _, err := ParseScheduledMeeting(info, termInfo)
				if err != nil {
					// didn't work, give up
					return nil, nil, nil, fmt.Errorf("mit: ParseScheduledMeeting: couldn't handle info '%s'", info)
				}

				// it worked, so keep going
				if parsedInfo.StartSeconds != 0 {
					scheduledMeeting.StartSeconds = parsedInfo.StartSeconds
				}
				if parsedInfo.EndSeconds != 0 {
					scheduledMeeting.EndSeconds = parsedInfo.EndSeconds
				}
			}
		}

		scheduledMeetingString = parensRegex.ReplaceAllString(scheduledMeetingString, "")
	}

	scheduledMeetingString = strings.TrimSpace(scheduledMeetingString)

	charMap := map[rune]time.Weekday{
		'M': time.Monday,
		'T': time.Tuesday,
		'W': time.Wednesday,
		'R': time.Thursday,
		'F': time.Friday,
		'S': time.Saturday,
	}

	remainingScheduledMeetingString := ""

	// first parse out all the weekdays
	for i, scheduledMeetingRune := range scheduledMeetingString {
		weekday, isWeekday := charMap[scheduledMeetingRune]
		if isWeekday {
			scheduledMeeting.Weekdays = append(scheduledMeeting.Weekdays, weekday)
		} else {
			// we're done
			remainingScheduledMeetingString = scheduledMeetingString[i:]
			break
		}
	}

	var err error
	if remainingScheduledMeetingString != "" && !parsedATime {
		if strings.Contains(remainingScheduledMeetingString, "-") {
			// it's a time like "4-5.30"
			forceAM := false
			if strings.Contains(remainingScheduledMeetingString, "AM") {
				forceAM = true
				remainingScheduledMeetingString = strings.Replace(remainingScheduledMeetingString, "AM", "", -1)
			}

			// HACK: 18.204 has a string in the format "MW2.30-4 (LIMITED 15 EACH S .."
			// we detect this and ignore the parenthesis
			if strings.Contains(remainingScheduledMeetingString, "(") && !strings.Contains(remainingScheduledMeetingString, ")") {
				remainingScheduledMeetingString = strings.Split(remainingScheduledMeetingString, "(")[0]
			}

			timeParts := strings.Split(remainingScheduledMeetingString, "-")
			if len(timeParts) != 2 {
				return nil, nil, nil, fmt.Errorf("mit: ParseScheduledMeeting: time info string '%s' had too many dashes", scheduledMeetingString)
			}

			scheduledMeeting.StartSeconds, err = parseTime(timeParts[0], forceAM)
			if err != nil {
				return nil, nil, nil, err
			}

			scheduledMeeting.EndSeconds, err = parseTime(timeParts[1], forceAM)
			if err != nil {
				return nil, nil, nil, err
			}
		} else {
			// it's a time like "4"
			scheduledMeeting.StartSeconds, err = parseTime(remainingScheduledMeetingString, false)
			if err != nil {
				return nil, nil, nil, err
			}

			// we can assume it's an hour long
			scheduledMeeting.EndSeconds = scheduledMeeting.StartSeconds + 60*60
		}
	} else {
		if strings.TrimSpace(remainingScheduledMeetingString) != "" && parsedATime {
			// there's something left, and it can't be a time...
			// complain
			return nil, nil, nil, fmt.Errorf("mit: unexpected extra data '%s' when parsing time string '%s'", remainingScheduledMeetingString, scheduledMeetingString)
		}
	}

	return &scheduledMeeting, beginsOn, endsOn, nil
}

// ParseTimeInfo parses out the time information from an input string in the format "MW9-11,F9".
func ParseTimeInfo(timeInfoString string, termInfo TermInfo) (TimeInfo, error) {
	timeInfo := TimeInfo{
		Meetings: []ScheduledMeeting{},
		BeginsOn: termInfo.FirstDayOfClasses,
		EndsOn:   termInfo.LastDayOfClasses,
	}

	// hack to cope with the fact that some people really, really don't understand how the course catalog works
	// this is specifically for AS.301
	timeInfoString = strings.Replace(timeInfoString, "Thurs", "R", -1)

	timeInfoParts := strings.Split(timeInfoString, ",")
	for _, part := range timeInfoParts {
		meeting, beginsOn, endsOn, err := ParseScheduledMeeting(part, termInfo)
		if err != nil {
			return TimeInfo{}, err
		}

		if beginsOn != nil {
			timeInfo.BeginsOn = *beginsOn
		}
		if endsOn != nil {
			timeInfo.EndsOn = *endsOn
		}

		if meeting != nil {
			timeInfo.Meetings = append(timeInfo.Meetings, *meeting)
		}
	}

	return timeInfo, nil
}
<commit_msg>mit/parse: add additional strings that indicate a subject is not scheduled on the calendar<commit_after>package mit

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var parensRegex = regexp.MustCompile("\\((.*?)\\)")

// A ScheduledMeeting struct contains details about a scheduled meeting in a TimeInfo. Each TimeInfo can have multiple ScheduledMeetings.
type ScheduledMeeting struct {
	Weekdays     []time.Weekday
	StartSeconds int
	EndSeconds   int
}

// A TimeInfo struct contains details about when a class's section meets.
type TimeInfo struct {
	Meetings []ScheduledMeeting
	BeginsOn time.Time
	EndsOn   time.Time
}

func parseTime(timeString string, forceAM bool) (int, error) {
	// sometimes, for whatever reason, use a colon instead of the normal dot as a separator
	// this seems to just be some advising seminars
	// replace these colons with a dot so that we can handle them
	timeString = strings.Replace(timeString, ":", ".", -1)

	// some classes also like to just randomly add a "pm"
	// looking at you WGS.228 and 15.389
	isPM := false
	if strings.Contains(timeString, "pm") || strings.Contains(timeString, "PM") {
		timeString = strings.Replace(strings.Replace(timeString, "pm", "", -1), "PM", "", -1)
		isPM = true
	}

	// also check for just "p"
	if strings.Contains(timeString, "p") || strings.Contains(timeString, "P") {
		timeString = strings.Replace(strings.Replace(timeString, "p", "", -1), "P", "", -1)
		isPM = true
	}

	// it's a time like "4" or "5.30"
	parts := strings.Split(timeString, ".")
	if len(parts) > 2 {
		return 0, fmt.Errorf("mit: parseTime: time string '%s' had too many dots", timeString)
	}

	hour, err := strconv.Atoi(strings.TrimSpace(parts[0]))
	if err != nil {
		return 0, err
	}

	minute := 0

	if len(parts) == 2 {
		minute, err = strconv.Atoi(strings.TrimSpace(parts[1]))
		if err != nil {
			return 0, err
		}
	}

	if forceAM || (!isPM && hour >= 8 && hour <= 12) {
		// assume AM
	} else {
		// assume PM
		hour += 12
		if hour == 24 {
			// 12 PM is actually noon
			hour = 12
		}
	}

	timeSeconds := ((hour * 60) + minute) * 60

	return timeSeconds, nil
}

func assignYearForDate(date time.Time, termInfo TermInfo) time.Time {
	return date.AddDate(termInfo.FirstDayOfClasses.Year(), 0, 0)
}

// ParseScheduledMeeting parses the given time info string, in a format like "MW4-5.30"
func ParseScheduledMeeting(scheduledMeetingString string, termInfo TermInfo) (*ScheduledMeeting, *time.Time, *time.Time, error) {
	scheduledMeeting := ScheduledMeeting{
		Weekdays: []time.Weekday{},
	}

	var beginsOn *time.Time
	var endsOn *time.Time

	scheduledMeetingString = strings.ToUpper(strings.TrimSpace(strings.Replace(scheduledMeetingString, "*", "", -1)))

	if scheduledMeetingString == "SUBJECT CANCELLED" || scheduledMeetingString == "SECTION CANCELLED" || scheduledMeetingString == "TO BE ARRANGED" || scheduledMeetingString == "TBA" || scheduledMeetingString == "TBD" || scheduledMeetingString == "ARRANGED" || scheduledMeetingString == "SCHEDULED INDIVIDUALLY" || scheduledMeetingString == "SCHECULED INDIVIDUALLY" || scheduledMeetingString == "UNSCHEDULED" || scheduledMeetingString == "UNSCEDULED" || scheduledMeetingString == "ARR" {
		// oof
		return nil, nil, nil, nil
	}

	// HACK: some classes use "TH" for thursday instead of a single character
	// for example, 12.803 has "TTH 10:30 - 12 PM"
	// we work around this by replacing "TH" with the correct code, "R"
	// this works because H is not used for anything else
	scheduledMeetingString = strings.ReplaceAll(scheduledMeetingString, "TH", "R")

	parsedATime := false

	if strings.Contains(scheduledMeetingString, "(") {
		// there's a thing in parentheses
		// for example "TR10.30-12 (BEGINS OCT 21)"

		// it could be an evening class
		// for example "TR EVE (4-6 PM)" or "W EVE (4-6.30 PM)"
		isEvening := (strings.Contains(scheduledMeetingString, "EVE"))

		matches := parensRegex.FindAllStringSubmatch(scheduledMeetingString, -1)
		for i, match := range matches {
			info := match[1]

			if i == 0 && isEvening {
				// it's an evening class
				// the first paren is a time
				info = strings.Replace(info, " PM", "", -1)

				subScheduledMeeting, _, _, err := ParseScheduledMeeting(info, termInfo)
				if err != nil {
					return nil, nil, nil, err
				}

				if subScheduledMeeting.StartSeconds < 12*60*60 {
					subScheduledMeeting.StartSeconds += 12 * 60 * 60
				} else if subScheduledMeeting.EndSeconds < 12*60*60 {
					subScheduledMeeting.EndSeconds += 12 * 60 * 60
				}

				scheduledMeeting.StartSeconds = subScheduledMeeting.StartSeconds
				scheduledMeeting.EndSeconds = subScheduledMeeting.EndSeconds

				scheduledMeetingString = strings.Replace(scheduledMeetingString, "EVE", "", -1)
				scheduledMeetingString = parensRegex.ReplaceAllString(scheduledMeetingString, "")

				parsedATime = true

				continue
			}

			if strings.HasPrefix(info, "BEGINS") {
				dateString := strings.Replace(info, "BEGINS ", "", -1)
				parsedDate, err := time.Parse("Jan 2", dateString)
				if err != nil {
					return nil, nil, nil, err
				}
				parsedDate = assignYearForDate(parsedDate, termInfo)
				beginsOn = &parsedDate
			} else if strings.HasPrefix(info, "END") {
				dateString := strings.Replace(info, "ENDS ", "", -1)
				parsedDate, err := time.Parse("Jan 2", dateString)
				if err != nil {
					return nil, nil, nil, err
				}
				parsedDate = assignYearForDate(parsedDate, termInfo)
				endsOn = &parsedDate
			} else if strings.HasPrefix(info, "MEETS") {
				// it's a message in the format "MEETS 4/7 TO 5/14"
				dateString := strings.Replace(info, "MEETS ", "", -1)
				parts := strings.Split(dateString, "TO")

				parsedBegin, err := time.Parse("1/2", strings.TrimSpace(parts[0]))
				if err != nil {
					return nil, nil, nil, err
				}
				parsedBegin = assignYearForDate(parsedBegin, termInfo)
				beginsOn = &parsedBegin

				parsedEnd, err := time.Parse("1/2", strings.TrimSpace(parts[1]))
				if err != nil {
					return nil, nil, nil, err
				}
				parsedEnd = assignYearForDate(parsedEnd, termInfo)
				endsOn = &parsedEnd
			} else if strings.HasPrefix(info, "LIMITED") {
				// it's about limited enrollment, which we don't care about
				// just ignore it
			} else {
				// some classes are not entered into the catalog correctly
				// they have a format like T(11-12:30)
				// if we can parse the time in parenthesis, assume that's the problem
				parsedInfo, _, _, err := ParseScheduledMeeting(info, termInfo)
				if err != nil {
					// didn't work, give up
					return nil, nil, nil, fmt.Errorf("mit: ParseScheduledMeeting: couldn't handle info '%s'", info)
				}

				// it worked, so keep going
				if parsedInfo.StartSeconds != 0 {
					scheduledMeeting.StartSeconds = parsedInfo.StartSeconds
				}
				if parsedInfo.EndSeconds != 0 {
					scheduledMeeting.EndSeconds = parsedInfo.EndSeconds
				}
			}
		}

		scheduledMeetingString = parensRegex.ReplaceAllString(scheduledMeetingString, "")
	}

	scheduledMeetingString = strings.TrimSpace(scheduledMeetingString)

	charMap := map[rune]time.Weekday{
		'M': time.Monday,
		'T': time.Tuesday,
		'W': time.Wednesday,
		'R': time.Thursday,
		'F': time.Friday,
		'S': time.Saturday,
	}

	remainingScheduledMeetingString := ""

	// first parse out all the weekdays
	for i, scheduledMeetingRune := range scheduledMeetingString {
		weekday, isWeekday := charMap[scheduledMeetingRune]
		if isWeekday {
			scheduledMeeting.Weekdays = append(scheduledMeeting.Weekdays, weekday)
		} else {
			// we're done
			remainingScheduledMeetingString = scheduledMeetingString[i:]
			break
		}
	}

	var err error
	if remainingScheduledMeetingString != "" && !parsedATime {
		if strings.Contains(remainingScheduledMeetingString, "-") {
			// it's a time like "4-5.30"
			forceAM := false
			if strings.Contains(remainingScheduledMeetingString, "AM") {
				forceAM = true
				remainingScheduledMeetingString = strings.Replace(remainingScheduledMeetingString, "AM", "", -1)
			}

			// HACK: 18.204 has a string in the format "MW2.30-4 (LIMITED 15 EACH S .."
			// we detect this and ignore the parenthesis
			if strings.Contains(remainingScheduledMeetingString, "(") && !strings.Contains(remainingScheduledMeetingString, ")") {
				remainingScheduledMeetingString = strings.Split(remainingScheduledMeetingString, "(")[0]
			}

			timeParts := strings.Split(remainingScheduledMeetingString, "-")
			if len(timeParts) != 2 {
				return nil, nil, nil, fmt.Errorf("mit: ParseScheduledMeeting: time info string '%s' had too many dashes", scheduledMeetingString)
			}

			scheduledMeeting.StartSeconds, err = parseTime(timeParts[0], forceAM)
			if err != nil {
				return nil, nil, nil, err
			}

			scheduledMeeting.EndSeconds, err = parseTime(timeParts[1], forceAM)
			if err != nil {
				return nil, nil, nil, err
			}
		} else {
			// it's a time like "4"
			scheduledMeeting.StartSeconds, err = parseTime(remainingScheduledMeetingString, false)
			if err != nil {
				return nil, nil, nil, err
			}

			// we can assume it's an hour long
			scheduledMeeting.EndSeconds = scheduledMeeting.StartSeconds + 60*60
		}
	} else {
		if strings.TrimSpace(remainingScheduledMeetingString) != "" && parsedATime {
			// there's something left, and it can't be a time...
			// complain
			return nil, nil, nil, fmt.Errorf("mit: unexpected extra data '%s' when parsing time string '%s'", remainingScheduledMeetingString, scheduledMeetingString)
		}
	}

	return &scheduledMeeting, beginsOn, endsOn, nil
}

// ParseTimeInfo parses out the time information from an input string in the format "MW9-11,F9".
func ParseTimeInfo(timeInfoString string, termInfo TermInfo) (TimeInfo, error) {
	timeInfo := TimeInfo{
		Meetings: []ScheduledMeeting{},
		BeginsOn: termInfo.FirstDayOfClasses,
		EndsOn:   termInfo.LastDayOfClasses,
	}

	// hack to cope with the fact that some people really, really don't understand how the course catalog works
	// this is specifically for AS.301
	timeInfoString = strings.Replace(timeInfoString, "Thurs", "R", -1)

	timeInfoParts := strings.Split(timeInfoString, ",")
	for _, part := range timeInfoParts {
		meeting, beginsOn, endsOn, err := ParseScheduledMeeting(part, termInfo)
		if err != nil {
			return TimeInfo{}, err
		}

		if beginsOn != nil {
			timeInfo.BeginsOn = *beginsOn
		}
		if endsOn != nil {
			timeInfo.EndsOn = *endsOn
		}

		if meeting != nil {
			timeInfo.Meetings = append(timeInfo.Meetings, *meeting)
		}
	}

	return timeInfo, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package applescript

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
)

const (
	script = `
on run argv
	set commandType to first item of argv as string
	if (commandType is equal to "add_track_to_playlist") then
		add_track_to_playlist(second item of argv, third item of argv)
	end if
	if (commandType is equal to "list_of_playlists") then
		list_of_playlists()
	end if
end run

on add_track_to_playlist(trackPath, playlistName)
	tell application "iTunes"
		add (trackPath as POSIX file) to playlist playlistName
	end tell
end add_track_to_playlist

on list_of_playlists()
	tell application "iTunes"
		get name of playlists
	end tell
end list_of_playlists
`
)

var scriptFile *os.File

func AddTrackToPlaylist(trackPath, playlistName string) error {
	_, err := executeOSAScript("add_track_to_playlist", trackPath, playlistName)
	return err
}

func ListOfPlaylists() (string, error) {
	return executeOSAScript("list_of_playlists")
}

// executeOSAScript executes AppleScript script with args and returns output and error.
func executeOSAScript(args ...string) (string, error) {
	if scriptFile == nil {
		var err error
		scriptFile, err = ioutil.TempFile("", "")
		if err != nil {
			return "", fmt.Errorf("couldn't create osascript file: %v", err)
		}
		if _, err = scriptFile.Write([]byte(script)); err != nil {
			return "", fmt.Errorf("couldn't write script to file: %v", err)
		}
	}

	out, err := exec.Command("osascript", append([]string{scriptFile.Name()}, args...)...).Output()
	return string(out), err
}
<commit_msg>applescript: Add logging of osascript<commit_after>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package applescript

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
)

const (
	script = `
on run argv
	set commandType to first item of argv as string
	if (commandType is equal to "add_track_to_playlist") then
		add_track_to_playlist(second item of argv, third item of argv)
	end if
	if (commandType is equal to "list_of_playlists") then
		list_of_playlists()
	end if
end run

on add_track_to_playlist(trackPath, playlistName)
	tell application "iTunes"
		add (trackPath as POSIX file) to playlist playlistName
	end tell
end add_track_to_playlist

on list_of_playlists()
	tell application "iTunes"
		get name of playlists
	end tell
end list_of_playlists
`
)

var scriptFile *os.File

func AddTrackToPlaylist(trackPath, playlistName string) error {
	_, err := executeOSAScript("add_track_to_playlist", trackPath, playlistName)
	return err
}

func ListOfPlaylists() (string, error) {
	return executeOSAScript("list_of_playlists")
}

// executeOSAScript executes AppleScript script with args and returns output and error.
func executeOSAScript(args ...string) (string, error) {
	if scriptFile == nil {
		var err error
		scriptFile, err = ioutil.TempFile("", "")
		if err != nil {
			return "", fmt.Errorf("couldn't create osascript file: %v", err)
		}
		if _, err = scriptFile.Write([]byte(script)); err != nil {
			return "", fmt.Errorf("couldn't write script to file: %v", err)
		}
	}

	args = append([]string{scriptFile.Name()}, args...)
	jww.INFO.Println("Executing osascript with args :", args)
	out, err := exec.Command("osascript", args...).Output()
	return string(out), err
}
<|endoftext|>


<commit_before>package util

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/glog"
)

var (
	client    *http.Client
	Transport *http.Transport
)

func init() {
	Transport = &http.Transport{
		MaxIdleConnsPerHost: 1024,
	}
	client = &http.Client{
		Transport: Transport,
	}
}

func PostBytes(url string, body []byte) ([]byte, error) {
	r, err := client.Post(url, "", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("Post to %s: %v", url, err)
	}
	defer r.Body.Close()
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, fmt.Errorf("Read response body: %v", err)
	}
	if r.StatusCode >= 400 {
		return nil, fmt.Errorf("%s: %s", url, r.Status)
	}
	return b, nil
}

func Post(url string, values url.Values) ([]byte, error) {
	r, err := client.PostForm(url, values)
	if err != nil {
		return nil, err
	}
	defer r.Body.Close()
	b, err := ioutil.ReadAll(r.Body)
	if r.StatusCode >= 400 {
		if err != nil {
			return nil, fmt.Errorf("%s: %d - %s", url, r.StatusCode, string(b))
		} else {
			return nil, fmt.Errorf("%s: %s", url, r.Status)
		}
	}
	if err != nil {
		return nil, err
	}
	return b, nil
}

//	github.com/chrislusf/seaweedfs/unmaintained/repeated_vacuum/repeated_vacuum.go
//	may need increasing http.Client.Timeout
func Get(url string) ([]byte, error) {
	r, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer r.Body.Close()
	b, err := ioutil.ReadAll(r.Body)
	if r.StatusCode >= 400 {
		return nil, fmt.Errorf("%s: %s", url, r.Status)
	}
	if err != nil {
		return nil, err
	}
	return b, nil
}

func Head(url string) (http.Header, error) {
	r, err := client.Head(url)
	if err != nil {
		return nil, err
	}
	defer CloseResponse(r)
	if r.StatusCode >= 400 {
		return nil, fmt.Errorf("%s: %s", url, r.Status)
	}
	return r.Header, nil
}

func Delete(url string, jwt string) error {
	req, err := http.NewRequest("DELETE", url, nil)
	if jwt != "" {
		req.Header.Set("Authorization", "BEARER "+string(jwt))
	}
	if err != nil {
		return err
	}
	resp, e := client.Do(req)
	if e != nil {
		return e
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	switch resp.StatusCode {
	case http.StatusNotFound, http.StatusAccepted, http.StatusOK:
		return nil
	}
	m := make(map[string]interface{})
	if e := json.Unmarshal(body, m); e == nil {
		if s, ok := m["error"].(string); ok {
			return errors.New(s)
		}
	}
	return errors.New(string(body))
}

func GetBufferStream(url string, values url.Values, allocatedBytes []byte, eachBuffer func([]byte)) error {
	r, err := client.PostForm(url, values)
	if err != nil {
		return err
	}
	defer CloseResponse(r)
	if r.StatusCode != 200 {
		return fmt.Errorf("%s: %s", url, r.Status)
	}
	for {
		n, err := r.Body.Read(allocatedBytes)
		if n > 0 {
			eachBuffer(allocatedBytes[:n])
		}
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
	}
}

func GetUrlStream(url string, values url.Values, readFn func(io.Reader) error) error {
	r, err := client.PostForm(url, values)
	if err != nil {
		return err
	}
	defer CloseResponse(r)
	if r.StatusCode != 200 {
		return fmt.Errorf("%s: %s", url, r.Status)
	}
	return readFn(r.Body)
}

func DownloadFile(fileUrl string) (filename string, header http.Header, rc io.ReadCloser, e error) {
	response, err := client.Get(fileUrl)
	if err != nil {
		return "", nil, nil, err
	}
	header = response.Header
	contentDisposition := response.Header["Content-Disposition"]
	if len(contentDisposition) > 0 {
		idx := strings.Index(contentDisposition[0], "filename=")
		if idx != -1 {
			filename = contentDisposition[0][idx+len("filename="):]
			filename = strings.Trim(filename, "\"")
		}
	}
	rc = response.Body
	return
}

func Do(req *http.Request) (resp *http.Response, err error) {
	return client.Do(req)
}

func NormalizeUrl(url string) string {
	if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
		return url
	}
	return "http://" + url
}

func ReadUrl(fileUrl string, cipherKey []byte, isGzipped bool, isFullChunk bool, offset int64, size int, buf []byte) (int64, error) {

	if cipherKey != nil {
		var n int
		err := readEncryptedUrl(fileUrl, cipherKey, isGzipped, isFullChunk, offset, size, func(data []byte) {
			n = copy(buf, data)
		})
		return int64(n), err
	}

	req, err := http.NewRequest("GET", fileUrl, nil)
	if err != nil {
		return 0, err
	}
	if !isFullChunk {
		req.Header.Add("Range", fmt.Sprintf("bytes=%d-%d", offset, offset+int64(size)-1))
	} else {
		req.Header.Set("Accept-Encoding", "gzip")
	}

	r, err := client.Do(req)
	if err != nil {
		return 0, err
	}

	defer r.Body.Close()
	if r.StatusCode >= 400 {
		return 0, fmt.Errorf("%s: %s", fileUrl, r.Status)
	}

	var reader io.ReadCloser
	contentEncoding := r.Header.Get("Content-Encoding")
	switch contentEncoding {
	case "gzip":
		reader, err = gzip.NewReader(r.Body)
		defer reader.Close()
	default:
		reader = r.Body
	}

	var (
		i, m int
		n    int64
	)

	// refers to https://github.com/golang/go/blob/master/src/bytes/buffer.go#L199
	// commit id c170b14c2c1cfb2fd853a37add92a82fd6eb4318
	for {
		m, err = reader.Read(buf[i:])
		i += m
		n += int64(m)
		if err == io.EOF {
			return n, nil
		}
		if err != nil {
			return n, err
		}
		if n == int64(len(buf)) {
			break
		}
	}
	// drains the response body to avoid memory leak
	data, _ := ioutil.ReadAll(reader)
	if len(data) != 0 {
		glog.V(1).Infof("%s reader has remaining %d bytes", contentEncoding, len(data))
	}
	return n, err
}

func ReadUrlAsStream(fileUrl string, cipherKey []byte, isContentGzipped bool, isFullChunk bool, offset int64, size int, fn func(data []byte)) error {

	if cipherKey != nil {
		return readEncryptedUrl(fileUrl, cipherKey, isContentGzipped, isFullChunk, offset, size, fn)
	}

	req, err := http.NewRequest("GET", fileUrl, nil)
	if err != nil {
		return err
	}

	if !isFullChunk {
		req.Header.Add("Range", fmt.Sprintf("bytes=%d-%d", offset, offset+int64(size)-1))
	}

	r, err := client.Do(req)
	if err != nil {
		return err
	}
	defer CloseResponse(r)
	if r.StatusCode >= 400 {
		return fmt.Errorf("%s: %s", fileUrl, r.Status)
	}

	var (
		m int
	)
	buf := make([]byte, 64*1024)

	for {
		m, err = r.Body.Read(buf)
		fn(buf[:m])
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
	}

}

func readEncryptedUrl(fileUrl string, cipherKey []byte, isContentGzipped bool, isFullChunk bool, offset int64, size int, fn func(data []byte)) error {
	encryptedData, err := Get(fileUrl)
	if err != nil {
		return fmt.Errorf("fetch %s: %v", fileUrl, err)
	}
	decryptedData, err := Decrypt(encryptedData, CipherKey(cipherKey))
	if err != nil {
		return fmt.Errorf("decrypt %s: %v", fileUrl, err)
	}
	if isContentGzipped {
		decryptedData, err = UnGzipData(decryptedData)
		if err != nil {
			return fmt.Errorf("unzip decrypt %s: %v", fileUrl, err)
		}
	}
	if len(decryptedData) < int(offset)+size {
		return fmt.Errorf("read decrypted %s size %d [%d, %d)", fileUrl, len(decryptedData), offset, int(offset)+size)
	}
	if isFullChunk {
		fn(decryptedData)
	} else {
		fn(decryptedData[int(offset) : int(offset)+size])
	}
	return nil
}

func ReadUrlAsReaderCloser(fileUrl string, rangeHeader string) (io.ReadCloser, error) {

	req, err := http.NewRequest("GET", fileUrl, nil)
	if err != nil {
		return nil, err
	}
	if rangeHeader != "" {
		req.Header.Add("Range", rangeHeader)
	}

	r, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	if r.StatusCode >= 400 {
		return nil, fmt.Errorf("%s: %s", fileUrl, r.Status)
	}

	return r.Body, nil
}

func CloseResponse(resp *http.Response) {
	io.Copy(ioutil.Discard, resp.Body)
	resp.Body.Close()
}
<commit_msg>fix "call of Unmarshal passes non-pointer as second argument"<commit_after>package util

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/glog"
)

var (
	client    *http.Client
	Transport *http.Transport
)

func init() {
	Transport = &http.Transport{
		MaxIdleConnsPerHost: 1024,
	}
	client = &http.Client{
		Transport: Transport,
	}
}

func PostBytes(url string, body []byte) ([]byte, error) {
	r, err := client.Post(url, "", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("Post to %s: %v", url, err)
	}
	defer r.Body.Close()
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, fmt.Errorf("Read response body: %v", err)
	}
	if r.StatusCode >= 400 {
		return nil, fmt.Errorf("%s: %s", url, r.Status)
	}
	return b, nil
}

func Post(url string, values url.Values) ([]byte, error) {
	r, err := client.PostForm(url, values)
	if err != nil {
		return nil, err
	}
	defer r.Body.Close()
	b, err := ioutil.ReadAll(r.Body)
	if r.StatusCode >= 400 {
		if err != nil {
			return nil, fmt.Errorf("%s: %d - %s", url, r.StatusCode, string(b))
		} else {
			return nil, fmt.Errorf("%s: %s", url, r.Status)
		}
	}
	if err != nil {
		return nil, err
	}
	return b, nil
}

//	github.com/chrislusf/seaweedfs/unmaintained/repeated_vacuum/repeated_vacuum.go
//	may need increasing http.Client.Timeout
func Get(url string) ([]byte, error) {
	r, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer r.Body.Close()
	b, err := ioutil.ReadAll(r.Body)
	if r.StatusCode >= 400 {
		return nil, fmt.Errorf("%s: %s", url, r.Status)
	}
	if err != nil {
		return nil, err
	}
	return b, nil
}

func Head(url string) (http.Header, error) {
	r, err := client.Head(url)
	if err != nil {
		return nil, err
	}
	defer CloseResponse(r)
	if r.StatusCode >= 400 {
		return nil, fmt.Errorf("%s: %s", url, r.Status)
	}
	return r.Header, nil
}

func Delete(url string, jwt string) error {
	req, err := http.NewRequest("DELETE", url, nil)
	if jwt != "" {
		req.Header.Set("Authorization", "BEARER "+string(jwt))
	}
	if err != nil {
		return err
	}
	resp, e := client.Do(req)
	if e != nil {
		return e
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	switch resp.StatusCode {
	case http.StatusNotFound, http.StatusAccepted, http.StatusOK:
		return nil
	}
	m := make(map[string]interface{})
	if e := json.Unmarshal(body, &m); e == nil {
		if s, ok := m["error"].(string); ok {
			return errors.New(s)
		}
	}
	return errors.New(string(body))
}

func GetBufferStream(url string, values url.Values, allocatedBytes []byte, eachBuffer func([]byte)) error {
	r, err := client.PostForm(url, values)
	if err != nil {
		return err
	}
	defer CloseResponse(r)
	if r.StatusCode != 200 {
		return fmt.Errorf("%s: %s", url, r.Status)
	}
	for {
		n, err := r.Body.Read(allocatedBytes)
		if n > 0 {
			eachBuffer(allocatedBytes[:n])
		}
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
	}
}

func GetUrlStream(url string, values url.Values, readFn func(io.Reader) error) error {
	r, err := client.PostForm(url, values)
	if err != nil {
		return err
	}
	defer CloseResponse(r)
	if r.StatusCode != 200 {
		return fmt.Errorf("%s: %s", url, r.Status)
	}
	return readFn(r.Body)
}

func DownloadFile(fileUrl string) (filename string, header http.Header, rc io.ReadCloser, e error) {
	response, err := client.Get(fileUrl)
	if err != nil {
		return "", nil, nil, err
	}
	header = response.Header
	contentDisposition := response.Header["Content-Disposition"]
	if len(contentDisposition) > 0 {
		idx := strings.Index(contentDisposition[0], "filename=")
		if idx != -1 {
			filename = contentDisposition[0][idx+len("filename="):]
			filename = strings.Trim(filename, "\"")
		}
	}
	rc = response.Body
	return
}

func Do(req *http.Request) (resp *http.Response, err error) {
	return client.Do(req)
}

func NormalizeUrl(url string) string {
	if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
		return url
	}
	return "http://" + url
}

func ReadUrl(fileUrl string, cipherKey []byte, isGzipped bool, isFullChunk bool, offset int64, size int, buf []byte) (int64, error) {

	if cipherKey != nil {
		var n int
		err := readEncryptedUrl(fileUrl, cipherKey, isGzipped, isFullChunk, offset, size, func(data []byte) {
			n = copy(buf, data)
		})
		return int64(n), err
	}

	req, err := http.NewRequest("GET", fileUrl, nil)
	if err != nil {
		return 0, err
	}
	if !isFullChunk {
		req.Header.Add("Range", fmt.Sprintf("bytes=%d-%d", offset, offset+int64(size)-1))
	} else {
		req.Header.Set("Accept-Encoding", "gzip")
	}

	r, err := client.Do(req)
	if err != nil {
		return 0, err
	}

	defer r.Body.Close()
	if r.StatusCode >= 400 {
		return 0, fmt.Errorf("%s: %s", fileUrl, r.Status)
	}

	var reader io.ReadCloser
	contentEncoding := r.Header.Get("Content-Encoding")
	switch contentEncoding {
	case "gzip":
		reader, err = gzip.NewReader(r.Body)
		defer reader.Close()
	default:
		reader = r.Body
	}

	var (
		i, m int
		n    int64
	)

	// refers to https://github.com/golang/go/blob/master/src/bytes/buffer.go#L199
	// commit id c170b14c2c1cfb2fd853a37add92a82fd6eb4318
	for {
		m, err = reader.Read(buf[i:])
		i += m
		n += int64(m)
		if err == io.EOF {
			return n, nil
		}
		if err != nil {
			return n, err
		}
		if n == int64(len(buf)) {
			break
		}
	}
	// drains the response body to avoid memory leak
	data, _ := ioutil.ReadAll(reader)
	if len(data) != 0 {
		glog.V(1).Infof("%s reader has remaining %d bytes", contentEncoding, len(data))
	}
	return n, err
}

func ReadUrlAsStream(fileUrl string, cipherKey []byte, isContentGzipped bool, isFullChunk bool, offset int64, size int, fn func(data []byte)) error {

	if cipherKey != nil {
		return readEncryptedUrl(fileUrl, cipherKey, isContentGzipped, isFullChunk, offset, size, fn)
	}

	req, err := http.NewRequest("GET", fileUrl, nil)
	if err != nil {
		return err
	}

	if !isFullChunk {
		req.Header.Add("Range", fmt.Sprintf("bytes=%d-%d", offset, offset+int64(size)-1))
	}

	r, err := client.Do(req)
	if err != nil {
		return err
	}
	defer CloseResponse(r)
	if r.StatusCode >= 400 {
		return fmt.Errorf("%s: %s", fileUrl, r.Status)
	}

	var (
		m int
	)
	buf := make([]byte, 64*1024)

	for {
		m, err = r.Body.Read(buf)
		fn(buf[:m])
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
	}

}

func readEncryptedUrl(fileUrl string, cipherKey []byte, isContentGzipped bool, isFullChunk bool, offset int64, size int, fn func(data []byte)) error {
	encryptedData, err := Get(fileUrl)
	if err != nil {
		return fmt.Errorf("fetch %s: %v", fileUrl, err)
	}
	decryptedData, err := Decrypt(encryptedData, CipherKey(cipherKey))
	if err != nil {
		return fmt.Errorf("decrypt %s: %v", fileUrl, err)
	}
	if isContentGzipped {
		decryptedData, err = UnGzipData(decryptedData)
		if err != nil {
			return fmt.Errorf("unzip decrypt %s: %v", fileUrl, err)
		}
	}
	if len(decryptedData) < int(offset)+size {
		return fmt.Errorf("read decrypted %s size %d [%d, %d)", fileUrl, len(decryptedData), offset, int(offset)+size)
	}
	if isFullChunk {
		fn(decryptedData)
	} else {
		fn(decryptedData[int(offset) : int(offset)+size])
	}
	return nil
}

func ReadUrlAsReaderCloser(fileUrl string, rangeHeader string) (io.ReadCloser, error) {

	req, err := http.NewRequest("GET", fileUrl, nil)
	if err != nil {
		return nil, err
	}
	if rangeHeader != "" {
		req.Header.Add("Range", rangeHeader)
	}

	r, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	if r.StatusCode >= 400 {
		return nil, fmt.Errorf("%s: %s", fileUrl, r.Status)
	}

	return r.Body, nil
}

func CloseResponse(resp *http.Response) {
	io.Copy(ioutil.Discard, resp.Body)
	resp.Body.Close()
}
<|endoftext|>


<commit_before>package copy

import (
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"strings"

	"github.com/containers/image/image"
	"github.com/containers/image/signature"
	"github.com/containers/image/transports"
	"github.com/containers/image/types"
)

// supportedDigests lists the supported blob digest types.
var supportedDigests = map[string]func() hash.Hash{
	"sha256": sha256.New,
}

type digestingReader struct {
	source           io.Reader
	digest           hash.Hash
	expectedDigest   []byte
	validationFailed bool
}

// newDigestingReader returns an io.Reader implementation with contents of source, which will eventually return a non-EOF error
// and set validationFailed to true if the source stream does not match expectedDigestString.
func newDigestingReader(source io.Reader, expectedDigestString string) (*digestingReader, error) {
	fields := strings.SplitN(expectedDigestString, ":", 2)
	if len(fields) != 2 {
		return nil, fmt.Errorf("Invalid digest specification %s", expectedDigestString)
	}
	fn, ok := supportedDigests[fields[0]]
	if !ok {
		return nil, fmt.Errorf("Invalid digest specification %s: unknown digest type %s", expectedDigestString, fields[0])
	}
	digest := fn()
	expectedDigest, err := hex.DecodeString(fields[1])
	if err != nil {
		return nil, fmt.Errorf("Invalid digest value %s: %v", expectedDigestString, err)
	}
	if len(expectedDigest) != digest.Size() {
		return nil, fmt.Errorf("Invalid digest specification %s: length %d does not match %d", expectedDigestString, len(expectedDigest), digest.Size())
	}
	return &digestingReader{
		source:           source,
		digest:           digest,
		expectedDigest:   expectedDigest,
		validationFailed: false,
	}, nil
}

func (d *digestingReader) Read(p []byte) (int, error) {
	n, err := d.source.Read(p)
	if n > 0 {
		if n2, err := d.digest.Write(p[:n]); n2 != n || err != nil {
			// Coverage: This should not happen, the hash.Hash interface requires
			// d.digest.Write to never return an error, and the io.Writer interface
			// requires n2 == len(input) if no error is returned.
			return 0, fmt.Errorf("Error updating digest during verification: %d vs. %d, %v", n2, n, err)
		}
	}
	if err == io.EOF {
		actualDigest := d.digest.Sum(nil)
		if subtle.ConstantTimeCompare(actualDigest, d.expectedDigest) != 1 {
			d.validationFailed = true
			return 0, fmt.Errorf("Digest did not match, expected %s, got %s", hex.EncodeToString(d.expectedDigest), hex.EncodeToString(actualDigest))
		}
	}
	return n, err
}

// Options allows supplying non-default configuration modifying the behavior of CopyImage.
type Options struct {
	RemoveSignatures bool   // Remove any pre-existing signatures. SignBy will still add a new signature.
	SignBy           string // If non-empty, asks for a signature to be added during the copy, and specifies a key ID, as accepted by signature.NewGPGSigningMechanism().SignDockerManifest(),
}

// Image copies image from srcRef to destRef, using policyContext to validate source image admissibility.
func Image(ctx *types.SystemContext, policyContext *signature.PolicyContext, destRef, srcRef types.ImageReference, options *Options) error {
	dest, err := destRef.NewImageDestination(ctx)
	if err != nil {
		return fmt.Errorf("Error initializing destination %s: %v", transports.ImageName(destRef), err)
	}
	defer dest.Close()

	rawSource, err := srcRef.NewImageSource(ctx, dest.SupportedManifestMIMETypes())
	if err != nil {
		return fmt.Errorf("Error initializing source %s: %v", transports.ImageName(srcRef), err)
	}
	src := image.FromSource(rawSource)
	defer src.Close()

	// Please keep this policy check BEFORE reading any other information about the image.
	if allowed, err := policyContext.IsRunningImageAllowed(src); !allowed || err != nil { // Be paranoid and fail if either return value indicates so.
		return fmt.Errorf("Source image rejected: %v", err)
	}

	manifest, _, err := src.Manifest()
	if err != nil {
		return fmt.Errorf("Error reading manifest: %v", err)
	}

	var sigs [][]byte
	if options != nil && options.RemoveSignatures {
		sigs = [][]byte{}
	} else {
		s, err := src.Signatures()
		if err != nil {
			return fmt.Errorf("Error reading signatures: %v", err)
		}
		sigs = s
	}
	if len(sigs) != 0 {
		if err := dest.SupportsSignatures(); err != nil {
			return fmt.Errorf("Can not copy signatures: %v", err)
		}
	}

	configInfo, err := src.ConfigInfo()
	if err != nil {
		return fmt.Errorf("Error parsing manifest: %v", err)
	}
	if configInfo.Digest != "" {
		if err := copyBlob(dest, rawSource, configInfo.Digest); err != nil {
			return err
		}
	}
	layerInfos, err := src.LayerInfos()
	if err != nil {
		return fmt.Errorf("Error parsing manifest: %v", err)
	}
	copiedLayers := map[string]struct{}{}
	for _, info := range layerInfos {
		if _, ok := copiedLayers[info.Digest]; !ok {
			if err := copyBlob(dest, rawSource, info.Digest); err != nil {
				return err
			}
			copiedLayers[info.Digest] = struct{}{}
		}
	}

	if options != nil && options.SignBy != "" {
		mech, err := signature.NewGPGSigningMechanism()
		if err != nil {
			return fmt.Errorf("Error initializing GPG: %v", err)
		}
		dockerReference := dest.Reference().DockerReference()
		if dockerReference == nil {
			return fmt.Errorf("Cannot determine canonical Docker reference for destination %s", transports.ImageName(dest.Reference()))
		}

		newSig, err := signature.SignDockerManifest(manifest, dockerReference.String(), mech, options.SignBy)
		if err != nil {
			return fmt.Errorf("Error creating signature: %v", err)
		}
		sigs = append(sigs, newSig)
	}

	if err := dest.PutManifest(manifest); err != nil {
		return fmt.Errorf("Error writing manifest: %v", err)
	}

	if err := dest.PutSignatures(sigs); err != nil {
		return fmt.Errorf("Error writing signatures: %v", err)
	}

	if err := dest.Commit(); err != nil {
		return fmt.Errorf("Error committing the finished image: %v", err)
	}

	return nil
}

func copyBlob(dest types.ImageDestination, src types.ImageSource, digest string) error {
	stream, blobSize, err := src.GetBlob(digest)
	if err != nil {
		return fmt.Errorf("Error reading blob %s: %v", digest, err)
	}
	defer stream.Close()

	// Be paranoid; in case PutBlob somehow managed to ignore an error from digestingReader,
	// use a separate validation failure indicator.
	// Note that we don't use a stronger "validationSucceeded" indicator, because
	// dest.PutBlob may detect that the layer already exists, in which case we don't
	// read stream to the end, and validation does not happen.
	digestingReader, err := newDigestingReader(stream, digest)
	if err != nil {
		return fmt.Errorf("Error preparing to verify blob %s: %v", digest, err)
	}
	if _, err := dest.PutBlob(digestingReader, types.BlobInfo{Digest: digest, Size: blobSize}); err != nil {
		return fmt.Errorf("Error writing blob: %v", err)
	}
	if digestingReader.validationFailed { // Coverage: This should never happen.
		return fmt.Errorf("Internal error uploading blob %s, digest verification failed but was ignored", digest)
	}
	return nil
}
<commit_msg>Add a minimal skeleton for manifest updates to copy.go<commit_after>package copy

import (
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"reflect"
	"strings"

	"github.com/containers/image/image"
	"github.com/containers/image/signature"
	"github.com/containers/image/transports"
	"github.com/containers/image/types"
)

// supportedDigests lists the supported blob digest types.
var supportedDigests = map[string]func() hash.Hash{
	"sha256": sha256.New,
}

type digestingReader struct {
	source           io.Reader
	digest           hash.Hash
	expectedDigest   []byte
	validationFailed bool
}

// newDigestingReader returns an io.Reader implementation with contents of source, which will eventually return a non-EOF error
// and set validationFailed to true if the source stream does not match expectedDigestString.
func newDigestingReader(source io.Reader, expectedDigestString string) (*digestingReader, error) {
	fields := strings.SplitN(expectedDigestString, ":", 2)
	if len(fields) != 2 {
		return nil, fmt.Errorf("Invalid digest specification %s", expectedDigestString)
	}
	fn, ok := supportedDigests[fields[0]]
	if !ok {
		return nil, fmt.Errorf("Invalid digest specification %s: unknown digest type %s", expectedDigestString, fields[0])
	}
	digest := fn()
	expectedDigest, err := hex.DecodeString(fields[1])
	if err != nil {
		return nil, fmt.Errorf("Invalid digest value %s: %v", expectedDigestString, err)
	}
	if len(expectedDigest) != digest.Size() {
		return nil, fmt.Errorf("Invalid digest specification %s: length %d does not match %d", expectedDigestString, len(expectedDigest), digest.Size())
	}
	return &digestingReader{
		source:           source,
		digest:           digest,
		expectedDigest:   expectedDigest,
		validationFailed: false,
	}, nil
}

func (d *digestingReader) Read(p []byte) (int, error) {
	n, err := d.source.Read(p)
	if n > 0 {
		if n2, err := d.digest.Write(p[:n]); n2 != n || err != nil {
			// Coverage: This should not happen, the hash.Hash interface requires
			// d.digest.Write to never return an error, and the io.Writer interface
			// requires n2 == len(input) if no error is returned.
			return 0, fmt.Errorf("Error updating digest during verification: %d vs. %d, %v", n2, n, err)
		}
	}
	if err == io.EOF {
		actualDigest := d.digest.Sum(nil)
		if subtle.ConstantTimeCompare(actualDigest, d.expectedDigest) != 1 {
			d.validationFailed = true
			return 0, fmt.Errorf("Digest did not match, expected %s, got %s", hex.EncodeToString(d.expectedDigest), hex.EncodeToString(actualDigest))
		}
	}
	return n, err
}

// Options allows supplying non-default configuration modifying the behavior of CopyImage.
type Options struct {
	RemoveSignatures bool   // Remove any pre-existing signatures. SignBy will still add a new signature.
	SignBy           string // If non-empty, asks for a signature to be added during the copy, and specifies a key ID, as accepted by signature.NewGPGSigningMechanism().SignDockerManifest(),
}

// Image copies image from srcRef to destRef, using policyContext to validate source image admissibility.
func Image(ctx *types.SystemContext, policyContext *signature.PolicyContext, destRef, srcRef types.ImageReference, options *Options) error {
	dest, err := destRef.NewImageDestination(ctx)
	if err != nil {
		return fmt.Errorf("Error initializing destination %s: %v", transports.ImageName(destRef), err)
	}
	defer dest.Close()

	rawSource, err := srcRef.NewImageSource(ctx, dest.SupportedManifestMIMETypes())
	if err != nil {
		return fmt.Errorf("Error initializing source %s: %v", transports.ImageName(srcRef), err)
	}
	src := image.FromSource(rawSource)
	defer src.Close()

	// Please keep this policy check BEFORE reading any other information about the image.
	if allowed, err := policyContext.IsRunningImageAllowed(src); !allowed || err != nil { // Be paranoid and fail if either return value indicates so.
		return fmt.Errorf("Source image rejected: %v", err)
	}

	manifest, _, err := src.Manifest()
	if err != nil {
		return fmt.Errorf("Error reading manifest: %v", err)
	}

	var sigs [][]byte
	if options != nil && options.RemoveSignatures {
		sigs = [][]byte{}
	} else {
		s, err := src.Signatures()
		if err != nil {
			return fmt.Errorf("Error reading signatures: %v", err)
		}
		sigs = s
	}
	if len(sigs) != 0 {
		if err := dest.SupportsSignatures(); err != nil {
			return fmt.Errorf("Can not copy signatures: %v", err)
		}
	}
	canModifyManifest := len(sigs) == 0

	configInfo, err := src.ConfigInfo()
	if err != nil {
		return fmt.Errorf("Error parsing manifest: %v", err)
	}
	if configInfo.Digest != "" {
		if err := copyBlob(dest, rawSource, configInfo.Digest); err != nil {
			return err
		}
	}
	layerInfos, err := src.LayerInfos()
	if err != nil {
		return fmt.Errorf("Error parsing manifest: %v", err)
	}
	copiedLayers := map[string]struct{}{}
	for _, info := range layerInfos {
		if _, ok := copiedLayers[info.Digest]; !ok {
			if err := copyBlob(dest, rawSource, info.Digest); err != nil {
				return err
			}
			copiedLayers[info.Digest] = struct{}{}
		}
	}

	manifestUpdates := types.ManifestUpdateOptions{}

	if !reflect.DeepEqual(manifestUpdates, types.ManifestUpdateOptions{}) {
		if !canModifyManifest {
			return fmt.Errorf("Internal error: copy needs an updated manifest but that was known to be forbidden")
		}
		manifest, err = src.UpdatedManifest(manifestUpdates)
		if err != nil {
			return fmt.Errorf("Error creating an updated manifest: %v", err)
		}
	}

	if options != nil && options.SignBy != "" {
		mech, err := signature.NewGPGSigningMechanism()
		if err != nil {
			return fmt.Errorf("Error initializing GPG: %v", err)
		}
		dockerReference := dest.Reference().DockerReference()
		if dockerReference == nil {
			return fmt.Errorf("Cannot determine canonical Docker reference for destination %s", transports.ImageName(dest.Reference()))
		}

		newSig, err := signature.SignDockerManifest(manifest, dockerReference.String(), mech, options.SignBy)
		if err != nil {
			return fmt.Errorf("Error creating signature: %v", err)
		}
		sigs = append(sigs, newSig)
	}

	if err := dest.PutManifest(manifest); err != nil {
		return fmt.Errorf("Error writing manifest: %v", err)
	}

	if err := dest.PutSignatures(sigs); err != nil {
		return fmt.Errorf("Error writing signatures: %v", err)
	}

	if err := dest.Commit(); err != nil {
		return fmt.Errorf("Error committing the finished image: %v", err)
	}

	return nil
}

func copyBlob(dest types.ImageDestination, src types.ImageSource, digest string) error {
	stream, blobSize, err := src.GetBlob(digest)
	if err != nil {
		return fmt.Errorf("Error reading blob %s: %v", digest, err)
	}
	defer stream.Close()

	// Be paranoid; in case PutBlob somehow managed to ignore an error from digestingReader,
	// use a separate validation failure indicator.
	// Note that we don't use a stronger "validationSucceeded" indicator, because
	// dest.PutBlob may detect that the layer already exists, in which case we don't
	// read stream to the end, and validation does not happen.
	digestingReader, err := newDigestingReader(stream, digest)
	if err != nil {
		return fmt.Errorf("Error preparing to verify blob %s: %v", digest, err)
	}
	if _, err := dest.PutBlob(digestingReader, types.BlobInfo{Digest: digest, Size: blobSize}); err != nil {
		return fmt.Errorf("Error writing blob: %v", err)
	}
	if digestingReader.validationFailed { // Coverage: This should never happen.
		return fmt.Errorf("Internal error uploading blob %s, digest verification failed but was ignored", digest)
	}
	return nil
}
<|endoftext|>


<commit_before>package pgx

import (
	"bytes"
	"context"
	"fmt"
	"io"

	"github.com/jackc/pgconn"
	"github.com/jackc/pgio"
	errors "golang.org/x/xerrors"
)

// CopyFromRows returns a CopyFromSource interface over the provided rows slice
// making it usable by *Conn.CopyFrom.
func CopyFromRows(rows [][]interface{}) CopyFromSource {
	return &copyFromRows{rows: rows, idx: -1}
}

type copyFromRows struct {
	rows [][]interface{}
	idx  int
}

func (ctr *copyFromRows) Next() bool {
	ctr.idx++
	return ctr.idx < len(ctr.rows)
}

func (ctr *copyFromRows) Values() ([]interface{}, error) {
	return ctr.rows[ctr.idx], nil
}

func (ctr *copyFromRows) Err() error {
	return nil
}

// CopyFromSlice returns a CopyFromSource interface over a dynamic func
// making it usable by *Conn.CopyFrom.
func CopyFromSlice(length int, next func(int) ([]interface{}, error)) CopyFromSource {
	return &copyFromSlice{next: next, idx: -1, len: length}
}

type copyFromSlice struct {
	next func(int) ([]interface{}, error)
	idx  int
	len  int
}

func (cts *copyFromSlice) Next() bool {
	cts.idx++
	return cts.idx < cts.len
}

func (cts *copyFromSlice) Values() ([]interface{}, error) {
	return cts.next(cts.idx)
}

func (cts *copyFromSlice) Err() error {
	return nil
}

// CopyFromSource is the interface used by *Conn.CopyFrom as the source for copy data.
type CopyFromSource interface {
	// Next returns true if there is another row and makes the next row data
	// available to Values(). When there are no more rows available or an error
	// has occurred it returns false.
	Next() bool

	// Values returns the values for the current row.
	Values() ([]interface{}, error)

	// Err returns any error that has been encountered by the CopyFromSource. If
	// this is not nil *Conn.CopyFrom will abort the copy.
	Err() error
}

type copyFrom struct {
	conn          *Conn
	tableName     Identifier
	columnNames   []string
	rowSrc        CopyFromSource
	readerErrChan chan error
}

func (ct *copyFrom) run(ctx context.Context) (int64, error) {
	quotedTableName := ct.tableName.Sanitize()
	cbuf := &bytes.Buffer{}
	for i, cn := range ct.columnNames {
		if i != 0 {
			cbuf.WriteString(", ")
		}
		cbuf.WriteString(quoteIdentifier(cn))
	}
	quotedColumnNames := cbuf.String()

	sd, err := ct.conn.Prepare(ctx, "", fmt.Sprintf("select %s from %s", quotedColumnNames, quotedTableName))
	if err != nil {
		return 0, err
	}

	r, w := io.Pipe()
	doneChan := make(chan struct{})

	go func() {
		defer close(doneChan)

		// Purposely NOT using defer w.Close(). See https://github.com/golang/go/issues/24283.
		buf := ct.conn.wbuf

		buf = append(buf, "PGCOPY\n\377\r\n\000"...)
		buf = pgio.AppendInt32(buf, 0)
		buf = pgio.AppendInt32(buf, 0)

		moreRows := true
		for moreRows {
			var err error
			moreRows, buf, err = ct.buildCopyBuf(buf, sd)
			if err != nil {
				w.CloseWithError(err)
				return
			}

			if ct.rowSrc.Err() != nil {
				w.CloseWithError(ct.rowSrc.Err())
				return
			}

			if len(buf) > 0 {
				_, err = w.Write(buf)
				if err != nil {
					w.Close()
					return
				}
			}

			buf = buf[:0]
		}

		w.Close()
	}()

	commandTag, err := ct.conn.pgConn.CopyFrom(ctx, r, fmt.Sprintf("copy %s ( %s ) from stdin binary;", quotedTableName, quotedColumnNames))

	r.Close()
	<-doneChan

	return commandTag.RowsAffected(), err
}

func (ct *copyFrom) buildCopyBuf(buf []byte, sd *pgconn.StatementDescription) (bool, []byte, error) {

	for ct.rowSrc.Next() {
		values, err := ct.rowSrc.Values()
		if err != nil {
			return false, nil, err
		}
		if len(values) != len(ct.columnNames) {
			return false, nil, errors.Errorf("expected %d values, got %d values", len(ct.columnNames), len(values))
		}

		buf = pgio.AppendInt16(buf, int16(len(ct.columnNames)))
		for i, val := range values {
			buf, err = encodePreparedStatementArgument(ct.conn.connInfo, buf, sd.Fields[i].DataTypeOID, val)
			if err != nil {
				return false, nil, err
			}
		}

		if len(buf) > 65536 {
			return true, buf, nil
		}
	}

	return false, buf, nil
}

// CopyFrom uses the PostgreSQL copy protocol to perform bulk data insertion.
// It returns the number of rows copied and an error.
//
// CopyFrom requires all values use the binary format. Almost all types
// implemented by pgx use the binary format by default. Types implementing
// Encoder can only be used if they encode to the binary format.
func (c *Conn) CopyFrom(ctx context.Context, tableName Identifier, columnNames []string, rowSrc CopyFromSource) (int64, error) {
	ct := &copyFrom{
		conn:          c,
		tableName:     tableName,
		columnNames:   columnNames,
		rowSrc:        rowSrc,
		readerErrChan: make(chan error),
	}

	return ct.run(ctx)
}
<commit_msg>CopyFromSlice should remember error<commit_after>package pgx

import (
	"bytes"
	"context"
	"fmt"
	"io"

	"github.com/jackc/pgconn"
	"github.com/jackc/pgio"
	errors "golang.org/x/xerrors"
)

// CopyFromRows returns a CopyFromSource interface over the provided rows slice
// making it usable by *Conn.CopyFrom.
func CopyFromRows(rows [][]interface{}) CopyFromSource {
	return &copyFromRows{rows: rows, idx: -1}
}

type copyFromRows struct {
	rows [][]interface{}
	idx  int
}

func (ctr *copyFromRows) Next() bool {
	ctr.idx++
	return ctr.idx < len(ctr.rows)
}

func (ctr *copyFromRows) Values() ([]interface{}, error) {
	return ctr.rows[ctr.idx], nil
}

func (ctr *copyFromRows) Err() error {
	return nil
}

// CopyFromSlice returns a CopyFromSource interface over a dynamic func
// making it usable by *Conn.CopyFrom.
func CopyFromSlice(length int, next func(int) ([]interface{}, error)) CopyFromSource {
	return &copyFromSlice{next: next, idx: -1, len: length}
}

type copyFromSlice struct {
	next func(int) ([]interface{}, error)
	idx  int
	len  int
	err  error
}

func (cts *copyFromSlice) Next() bool {
	cts.idx++
	return cts.idx < cts.len
}

func (cts *copyFromSlice) Values() ([]interface{}, error) {
	values, err := cts.next(cts.idx)
	if err != nil {
		cts.err = err
	}
	return values, err
}

func (cts *copyFromSlice) Err() error {
	return cts.err
}

// CopyFromSource is the interface used by *Conn.CopyFrom as the source for copy data.
type CopyFromSource interface {
	// Next returns true if there is another row and makes the next row data
	// available to Values(). When there are no more rows available or an error
	// has occurred it returns false.
	Next() bool

	// Values returns the values for the current row.
	Values() ([]interface{}, error)

	// Err returns any error that has been encountered by the CopyFromSource. If
	// this is not nil *Conn.CopyFrom will abort the copy.
	Err() error
}

type copyFrom struct {
	conn          *Conn
	tableName     Identifier
	columnNames   []string
	rowSrc        CopyFromSource
	readerErrChan chan error
}

func (ct *copyFrom) run(ctx context.Context) (int64, error) {
	quotedTableName := ct.tableName.Sanitize()
	cbuf := &bytes.Buffer{}
	for i, cn := range ct.columnNames {
		if i != 0 {
			cbuf.WriteString(", ")
		}
		cbuf.WriteString(quoteIdentifier(cn))
	}
	quotedColumnNames := cbuf.String()

	sd, err := ct.conn.Prepare(ctx, "", fmt.Sprintf("select %s from %s", quotedColumnNames, quotedTableName))
	if err != nil {
		return 0, err
	}

	r, w := io.Pipe()
	doneChan := make(chan struct{})

	go func() {
		defer close(doneChan)

		// Purposely NOT using defer w.Close(). See https://github.com/golang/go/issues/24283.
		buf := ct.conn.wbuf

		buf = append(buf, "PGCOPY\n\377\r\n\000"...)
		buf = pgio.AppendInt32(buf, 0)
		buf = pgio.AppendInt32(buf, 0)

		moreRows := true
		for moreRows {
			var err error
			moreRows, buf, err = ct.buildCopyBuf(buf, sd)
			if err != nil {
				w.CloseWithError(err)
				return
			}

			if ct.rowSrc.Err() != nil {
				w.CloseWithError(ct.rowSrc.Err())
				return
			}

			if len(buf) > 0 {
				_, err = w.Write(buf)
				if err != nil {
					w.Close()
					return
				}
			}

			buf = buf[:0]
		}

		w.Close()
	}()

	commandTag, err := ct.conn.pgConn.CopyFrom(ctx, r, fmt.Sprintf("copy %s ( %s ) from stdin binary;", quotedTableName, quotedColumnNames))

	r.Close()
	<-doneChan

	return commandTag.RowsAffected(), err
}

func (ct *copyFrom) buildCopyBuf(buf []byte, sd *pgconn.StatementDescription) (bool, []byte, error) {

	for ct.rowSrc.Next() {
		values, err := ct.rowSrc.Values()
		if err != nil {
			return false, nil, err
		}
		if len(values) != len(ct.columnNames) {
			return false, nil, errors.Errorf("expected %d values, got %d values", len(ct.columnNames), len(values))
		}

		buf = pgio.AppendInt16(buf, int16(len(ct.columnNames)))
		for i, val := range values {
			buf, err = encodePreparedStatementArgument(ct.conn.connInfo, buf, sd.Fields[i].DataTypeOID, val)
			if err != nil {
				return false, nil, err
			}
		}

		if len(buf) > 65536 {
			return true, buf, nil
		}
	}

	return false, buf, nil
}

// CopyFrom uses the PostgreSQL copy protocol to perform bulk data insertion.
// It returns the number of rows copied and an error.
//
// CopyFrom requires all values use the binary format. Almost all types
// implemented by pgx use the binary format by default. Types implementing
// Encoder can only be used if they encode to the binary format.
func (c *Conn) CopyFrom(ctx context.Context, tableName Identifier, columnNames []string, rowSrc CopyFromSource) (int64, error) {
	ct := &copyFrom{
		conn:          c,
		tableName:     tableName,
		columnNames:   columnNames,
		rowSrc:        rowSrc,
		readerErrChan: make(chan error),
	}

	return ct.run(ctx)
}
<|endoftext|>