<commit_before>// Copyright 2013 Michael Yang. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
package v2

import (
	"errors"
	"fmt"
	"github.com/mikkyang/id3-go/encodedbytes"
)

const (
	FrameHeaderSize = 10
)

// FrameType holds frame id metadata and constructor method
// A set number of these are created in the version specific files
type FrameType struct {
	id          string
	description string
	constructor func(FrameHead, []byte) Framer
}

// Framer provides a generic interface for frames
// This is the default type returned when creating frames
type Framer interface {
	Id() string
	Size() uint
	StatusFlags() byte
	FormatFlags() byte
	String() string
	Bytes() []byte
	setOwner(*Tag)
}

// FrameHead represents the header of each frame
// Additional metadata is kept through the embedded frame type
// These do not usually need to be manually created
type FrameHead struct {
	FrameType
	statusFlags byte
	formatFlags byte
	size        uint32
	owner       *Tag
}

func (ft FrameType) Id() string {
	return ft.id
}

func (h FrameHead) Size() uint {
	return uint(h.size)
}

func (h *FrameHead) changeSize(diff int) {
	if diff >= 0 {
		h.size += uint32(diff)
	} else {
		h.size -= uint32(-diff)
	}

	if h.owner != nil {
		h.owner.changeSize(diff)
	}
}

func (h FrameHead) StatusFlags() byte {
	return h.statusFlags
}

func (h FrameHead) FormatFlags() byte {
	return h.formatFlags
}

func (h *FrameHead) setOwner(t *Tag) {
	h.owner = t
}

// DataFrame is the default frame for binary data
type DataFrame struct {
	FrameHead
	data []byte
}

func NewDataFrame(ft FrameType, data []byte) *DataFrame {
	head := FrameHead{
		FrameType: ft,
		size:      uint32(len(data)),
	}

	return &DataFrame{head, data}
}

func ParseDataFrame(head FrameHead, data []byte) Framer {
	return &DataFrame{head, data}
}

func (f DataFrame) Data() []byte {
	return f.data
}

func (f *DataFrame) SetData(b []byte) {
	diff := len(b) - len(f.data)
	f.changeSize(diff)
	f.data = b
}

func (f DataFrame) String() string {
	return "<binary data>"
}

func (f DataFrame) Bytes() []byte {
	return f.data
}

// IdFrame represents identification tags
type IdFrame struct {
	FrameHead
	ownerIdentifier string
	identifier      []byte
}

func NewIdFrame(ft FrameType, ownerId string, id []byte) *IdFrame {
	head := FrameHead{
		FrameType: ft,
		size:      uint32(1 + len(ownerId) + len(id)),
	}

	return &IdFrame{
		FrameHead:       head,
		ownerIdentifier: ownerId,
		identifier:      id,
	}
}

func ParseIdFrame(head FrameHead, data []byte) Framer {
	var err error
	f := &IdFrame{FrameHead: head}
	rd := encodedbytes.NewReader(data)

	if f.ownerIdentifier, err = rd.ReadNullTermString(encodedbytes.NativeEncoding); err != nil {
		return nil
	}

	if f.identifier, err = rd.ReadRest(); len(f.identifier) > 64 || err != nil {
		return nil
	}

	return f
}

func (f IdFrame) OwnerIdentifier() string {
	return f.ownerIdentifier
}

func (f *IdFrame) SetOwnerIdentifier(ownerId string) {
	f.changeSize(len(ownerId) - len(f.ownerIdentifier))
	f.ownerIdentifier = ownerId
}

func (f IdFrame) Identifier() []byte {
	return f.identifier
}

func (f *IdFrame) SetIdentifier(id []byte) error {
	if len(id) > 64 {
		return errors.New("identifier: identifier too long")
	}

	f.changeSize(len(id) - len(f.identifier))
	f.identifier = id

	return nil
}

func (f IdFrame) String() string {
	return fmt.Sprintf("%s: %v", f.ownerIdentifier, f.identifier)
}

func (f IdFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteString(f.ownerIdentifier, encodedbytes.NativeEncoding); err != nil {
		return bytes
	}

	if _, err = wr.Write(f.identifier); err != nil {
		return bytes
	}

	return bytes
}

// TextFramer represents frames that contain encoded text
type TextFramer interface {
	Framer
	Encoding() string
	SetEncoding(string) error
	Text() string
	SetText(string) error
}

// TextFrame represents frames that contain encoded text
type TextFrame struct {
	FrameHead
	encoding byte
	text     string
}

func NewTextFrame(ft FrameType, text string) *TextFrame {
	head := FrameHead{
		FrameType: ft,
		size:      uint32(1 + len(text)),
	}

	return &TextFrame{
		FrameHead: head,
		text:      text,
	}
}

func ParseTextFrame(head FrameHead, data []byte) Framer {
	var err error
	f := &TextFrame{FrameHead: head}
	rd := encodedbytes.NewReader(data)

	if f.encoding, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.text, err = rd.ReadRestString(f.encoding); err != nil {
		return nil
	}

	return f
}

func (f TextFrame) Encoding() string {
	return encodedbytes.EncodingForIndex(f.encoding)
}

func (f *TextFrame) SetEncoding(encoding string) error {
	i := byte(encodedbytes.IndexForEncoding(encoding))
	if i < 0 {
		return errors.New("encoding: invalid encoding")
	}

	diff, err := encodedbytes.EncodedDiff(i, f.text, f.encoding, f.text)
	if err != nil {
		return err
	}

	f.changeSize(diff)
	f.encoding = i
	return nil
}

func (f TextFrame) Text() string {
	return f.text
}

func (f *TextFrame) SetText(text string) error {
	diff, err := encodedbytes.EncodedDiff(f.encoding, text, f.encoding, f.text)
	if err != nil {
		return err
	}

	f.changeSize(diff)
	f.text = text
	return nil
}

func (f TextFrame) String() string {
	return f.text
}

func (f TextFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteByte(f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.text, f.encoding); err != nil {
		return bytes
	}

	return bytes
}

type DescTextFrame struct {
	TextFrame
	description string
}

func NewDescTextFrame(ft FrameType, desc, text string) *DescTextFrame {
	f := NewTextFrame(ft, text)
	f.size += uint32(len(desc))

	return &DescTextFrame{
		TextFrame:   *f,
		description: desc,
	}
}

// DescTextFrame represents frames that contain encoded text and descriptions
func ParseDescTextFrame(head FrameHead, data []byte) Framer {
	var err error
	f := new(DescTextFrame)
	f.FrameHead = head
	rd := encodedbytes.NewReader(data)

	if f.encoding, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.description, err = rd.ReadNullTermString(f.encoding); err != nil {
		return nil
	}

	if f.text, err = rd.ReadRestString(f.encoding); err != nil {
		return nil
	}

	return f
}

func (f DescTextFrame) Description() string {
	return f.description
}

func (f *DescTextFrame) SetDescription(description string) error {
	diff, err := encodedbytes.EncodedDiff(f.encoding, description, f.encoding, f.description)
	if err != nil {
		return err
	}

	f.changeSize(diff)
	f.description = description
	return nil
}

func (f *DescTextFrame) SetEncoding(encoding string) error {
	i := byte(encodedbytes.IndexForEncoding(encoding))
	if i < 0 {
		return errors.New("encoding: invalid encoding")
	}

	descDiff, err := encodedbytes.EncodedDiff(i, f.text, f.encoding, f.text)
	if err != nil {
		return err
	}

	textDiff, err := encodedbytes.EncodedDiff(i, f.description, f.encoding, f.description)
	if err != nil {
		return err
	}

	f.changeSize(descDiff + textDiff)
	f.encoding = i
	return nil
}

func (f DescTextFrame) String() string {
	return fmt.Sprintf("%s: %s", f.description, f.text)
}

func (f DescTextFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteByte(f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.description, f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.text, f.encoding); err != nil {
		return bytes
	}

	return bytes
}

// UnsynchTextFrame represents frames that contain unsynchronized text
type UnsynchTextFrame struct {
	DescTextFrame
	language string
}

func NewUnsynchTextFrame(ft FrameType, desc, text string) *UnsynchTextFrame {
	f := NewDescTextFrame(ft, desc, text)
	f.size += uint32(3)

	// add null length for this encoding
	f.size += uint32(encodedbytes.EncodingNullLengthForIndex(f.encoding))

	return &UnsynchTextFrame{
		DescTextFrame: *f,
		language:      "eng",
	}
}

func ParseUnsynchTextFrame(head FrameHead, data []byte) Framer {
	var err error
	f := new(UnsynchTextFrame)
	f.FrameHead = head
	rd := encodedbytes.NewReader(data)

	if f.encoding, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.language, err = rd.ReadNumBytesString(3); err != nil {
		return nil
	}

	if f.description, err = rd.ReadNullTermString(f.encoding); err != nil {
		return nil
	}

	if f.text, err = rd.ReadRestString(f.encoding); err != nil {
		return nil
	}

	return f
}

func (f UnsynchTextFrame) Language() string {
	return f.language
}

func (f *UnsynchTextFrame) SetLanguage(language string) error {
	if len(language) != 3 {
		return errors.New("language: invalid language string")
	}

	f.language = language
	f.changeSize(0)
	return nil
}

func (f *UnsynchTextFrame) SetEncoding(encoding string) error {
	prevIndex := f.encoding
	err := f.DescTextFrame.SetEncoding(encoding)
	if err != nil {
		return err
	}

	n1 := encodedbytes.EncodingNullLengthForIndex(prevIndex)
	n2 := encodedbytes.EncodingNullLengthForIndex(f.encoding)
	f.changeSize(n2 - n1)
	return nil
}

func (f UnsynchTextFrame) String() string {
	return fmt.Sprintf("%s\t%s:\n%s", f.language, f.description, f.text)
}

func (f UnsynchTextFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteByte(f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.language, encodedbytes.NativeEncoding); err != nil {
		return bytes
	}

	if err = wr.WriteNullTermString(f.description, f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.text, f.encoding); err != nil {
		return bytes
	}

	return bytes
}

// ImageFrame represent frames that have media attached
type ImageFrame struct {
	DataFrame
	encoding    byte
	mimeType    string
	pictureType byte
	description string
}

func ParseImageFrame(head FrameHead, data []byte) Framer {
	var err error
	f := new(ImageFrame)
	f.FrameHead = head
	rd := encodedbytes.NewReader(data)

	if f.encoding, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.mimeType, err = rd.ReadNullTermString(encodedbytes.NativeEncoding); err != nil {
		return nil
	}

	if f.pictureType, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.description, err = rd.ReadNullTermString(f.encoding); err != nil {
		return nil
	}

	if f.data, err = rd.ReadRest(); err != nil {
		return nil
	}

	return f
}

func (f ImageFrame) Encoding() string {
	return encodedbytes.EncodingForIndex(f.encoding)
}

func (f *ImageFrame) SetEncoding(encoding string) error {
	i := byte(encodedbytes.IndexForEncoding(encoding))
	if i < 0 {
		return errors.New("encoding: invalid encoding")
	}

	diff, err := encodedbytes.EncodedDiff(i, f.description, f.encoding, f.description)
	if err != nil {
		return err
	}

	f.changeSize(diff)
	f.encoding = i
	return nil
}

func (f ImageFrame) MIMEType() string {
	return f.mimeType
}

func (f *ImageFrame) SetMIMEType(mimeType string) {
	diff := len(mimeType) - len(f.mimeType)
	if mimeType[len(mimeType)-1] != 0 {
		nullTermBytes := append([]byte(mimeType), 0x00)
		f.mimeType = string(nullTermBytes)
		diff += 1
	} else {
		f.mimeType = mimeType
	}

	f.changeSize(diff)
}

func (f ImageFrame) String() string {
	return fmt.Sprintf("%s\t%s: <binary data>", f.mimeType, f.description)
}

func (f ImageFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteByte(f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.mimeType, encodedbytes.NativeEncoding); err != nil {
		return bytes
	}

	if err = wr.WriteByte(f.pictureType); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.description, f.encoding); err != nil {
		return bytes
	}

	if n, err := wr.Write(f.data); n < len(f.data) || err != nil {
		return bytes
	}

	return bytes
}
<commit_msg>Fix sizes affected by encoding<commit_after>// Copyright 2013 Michael Yang. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
package v2

import (
	"errors"
	"fmt"
	"github.com/mikkyang/id3-go/encodedbytes"
)

const (
	FrameHeaderSize = 10
)

// FrameType holds frame id metadata and constructor method
// A set number of these are created in the version specific files
type FrameType struct {
	id          string
	description string
	constructor func(FrameHead, []byte) Framer
}

// Framer provides a generic interface for frames
// This is the default type returned when creating frames
type Framer interface {
	Id() string
	Size() uint
	StatusFlags() byte
	FormatFlags() byte
	String() string
	Bytes() []byte
	setOwner(*Tag)
}

// FrameHead represents the header of each frame
// Additional metadata is kept through the embedded frame type
// These do not usually need to be manually created
type FrameHead struct {
	FrameType
	statusFlags byte
	formatFlags byte
	size        uint32
	owner       *Tag
}

func (ft FrameType) Id() string {
	return ft.id
}

func (h FrameHead) Size() uint {
	return uint(h.size)
}

func (h *FrameHead) changeSize(diff int) {
	if diff >= 0 {
		h.size += uint32(diff)
	} else {
		h.size -= uint32(-diff)
	}

	if h.owner != nil {
		h.owner.changeSize(diff)
	}
}

func (h FrameHead) StatusFlags() byte {
	return h.statusFlags
}

func (h FrameHead) FormatFlags() byte {
	return h.formatFlags
}

func (h *FrameHead) setOwner(t *Tag) {
	h.owner = t
}

// DataFrame is the default frame for binary data
type DataFrame struct {
	FrameHead
	data []byte
}

func NewDataFrame(ft FrameType, data []byte) *DataFrame {
	head := FrameHead{
		FrameType: ft,
		size:      uint32(len(data)),
	}

	return &DataFrame{head, data}
}

func ParseDataFrame(head FrameHead, data []byte) Framer {
	return &DataFrame{head, data}
}

func (f DataFrame) Data() []byte {
	return f.data
}

func (f *DataFrame) SetData(b []byte) {
	diff := len(b) - len(f.data)
	f.changeSize(diff)
	f.data = b
}

func (f DataFrame) String() string {
	return "<binary data>"
}

func (f DataFrame) Bytes() []byte {
	return f.data
}

// IdFrame represents identification tags
type IdFrame struct {
	FrameHead
	ownerIdentifier string
	identifier      []byte
}

func NewIdFrame(ft FrameType, ownerId string, id []byte) *IdFrame {
	head := FrameHead{
		FrameType: ft,
		size:      uint32(1 + len(ownerId) + len(id)),
	}

	return &IdFrame{
		FrameHead:       head,
		ownerIdentifier: ownerId,
		identifier:      id,
	}
}

func ParseIdFrame(head FrameHead, data []byte) Framer {
	var err error
	f := &IdFrame{FrameHead: head}
	rd := encodedbytes.NewReader(data)

	if f.ownerIdentifier, err = rd.ReadNullTermString(encodedbytes.NativeEncoding); err != nil {
		return nil
	}

	if f.identifier, err = rd.ReadRest(); len(f.identifier) > 64 || err != nil {
		return nil
	}

	return f
}

func (f IdFrame) OwnerIdentifier() string {
	return f.ownerIdentifier
}

func (f *IdFrame) SetOwnerIdentifier(ownerId string) {
	f.changeSize(len(ownerId) - len(f.ownerIdentifier))
	f.ownerIdentifier = ownerId
}

func (f IdFrame) Identifier() []byte {
	return f.identifier
}

func (f *IdFrame) SetIdentifier(id []byte) error {
	if len(id) > 64 {
		return errors.New("identifier: identifier too long")
	}

	f.changeSize(len(id) - len(f.identifier))
	f.identifier = id

	return nil
}

func (f IdFrame) String() string {
	return fmt.Sprintf("%s: %v", f.ownerIdentifier, f.identifier)
}

func (f IdFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteString(f.ownerIdentifier, encodedbytes.NativeEncoding); err != nil {
		return bytes
	}

	if _, err = wr.Write(f.identifier); err != nil {
		return bytes
	}

	return bytes
}

// TextFramer represents frames that contain encoded text
type TextFramer interface {
	Framer
	Encoding() string
	SetEncoding(string) error
	Text() string
	SetText(string) error
}

// TextFrame represents frames that contain encoded text
type TextFrame struct {
	FrameHead
	encoding byte
	text     string
}

func NewTextFrame(ft FrameType, text string) *TextFrame {
	head := FrameHead{
		FrameType: ft,
		size:      uint32(1 + len(text)),
	}

	return &TextFrame{
		FrameHead: head,
		text:      text,
	}
}

func ParseTextFrame(head FrameHead, data []byte) Framer {
	var err error
	f := &TextFrame{FrameHead: head}
	rd := encodedbytes.NewReader(data)

	if f.encoding, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.text, err = rd.ReadRestString(f.encoding); err != nil {
		return nil
	}

	return f
}

func (f TextFrame) Encoding() string {
	return encodedbytes.EncodingForIndex(f.encoding)
}

func (f *TextFrame) SetEncoding(encoding string) error {
	i := byte(encodedbytes.IndexForEncoding(encoding))
	if i < 0 {
		return errors.New("encoding: invalid encoding")
	}

	diff, err := encodedbytes.EncodedDiff(i, f.text, f.encoding, f.text)
	if err != nil {
		return err
	}

	f.changeSize(diff)
	f.encoding = i
	return nil
}

func (f TextFrame) Text() string {
	return f.text
}

func (f *TextFrame) SetText(text string) error {
	diff, err := encodedbytes.EncodedDiff(f.encoding, text, f.encoding, f.text)
	if err != nil {
		return err
	}

	f.changeSize(diff)
	f.text = text
	return nil
}

func (f TextFrame) String() string {
	return f.text
}

func (f TextFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteByte(f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.text, f.encoding); err != nil {
		return bytes
	}

	return bytes
}

type DescTextFrame struct {
	TextFrame
	description string
}

func NewDescTextFrame(ft FrameType, desc, text string) *DescTextFrame {
	f := NewTextFrame(ft, text)
	nullLength := encodedbytes.EncodingNullLengthForIndex(f.encoding)
	f.size += uint32(len(desc) + nullLength)

	return &DescTextFrame{
		TextFrame:   *f,
		description: desc,
	}
}

// DescTextFrame represents frames that contain encoded text and descriptions
func ParseDescTextFrame(head FrameHead, data []byte) Framer {
	var err error
	f := new(DescTextFrame)
	f.FrameHead = head
	rd := encodedbytes.NewReader(data)

	if f.encoding, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.description, err = rd.ReadNullTermString(f.encoding); err != nil {
		return nil
	}

	if f.text, err = rd.ReadRestString(f.encoding); err != nil {
		return nil
	}

	return f
}

func (f DescTextFrame) Description() string {
	return f.description
}

func (f *DescTextFrame) SetDescription(description string) error {
	diff, err := encodedbytes.EncodedDiff(f.encoding, description, f.encoding, f.description)
	if err != nil {
		return err
	}

	f.changeSize(diff)
	f.description = description
	return nil
}

func (f *DescTextFrame) SetEncoding(encoding string) error {
	i := byte(encodedbytes.IndexForEncoding(encoding))
	if i < 0 {
		return errors.New("encoding: invalid encoding")
	}

	descDiff, err := encodedbytes.EncodedDiff(i, f.text, f.encoding, f.text)
	if err != nil {
		return err
	}

	newNullLength := encodedbytes.EncodingNullLengthForIndex(i)
	oldNullLength := encodedbytes.EncodingNullLengthForIndex(f.encoding)
	nullDiff := newNullLength - oldNullLength

	textDiff, err := encodedbytes.EncodedDiff(i, f.description, f.encoding, f.description)
	if err != nil {
		return err
	}

	f.changeSize(descDiff + nullDiff + textDiff)
	f.encoding = i
	return nil
}

func (f DescTextFrame) String() string {
	return fmt.Sprintf("%s: %s", f.description, f.text)
}

func (f DescTextFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteByte(f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.description, f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.text, f.encoding); err != nil {
		return bytes
	}

	return bytes
}

// UnsynchTextFrame represents frames that contain unsynchronized text
type UnsynchTextFrame struct {
	DescTextFrame
	language string
}

func NewUnsynchTextFrame(ft FrameType, desc, text string) *UnsynchTextFrame {
	f := NewDescTextFrame(ft, desc, text)
	f.size += uint32(3)

	return &UnsynchTextFrame{
		DescTextFrame: *f,
		language:      "eng",
	}
}

func ParseUnsynchTextFrame(head FrameHead, data []byte) Framer {
	var err error
	f := new(UnsynchTextFrame)
	f.FrameHead = head
	rd := encodedbytes.NewReader(data)

	if f.encoding, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.language, err = rd.ReadNumBytesString(3); err != nil {
		return nil
	}

	if f.description, err = rd.ReadNullTermString(f.encoding); err != nil {
		return nil
	}

	if f.text, err = rd.ReadRestString(f.encoding); err != nil {
		return nil
	}

	return f
}

func (f UnsynchTextFrame) Language() string {
	return f.language
}

func (f *UnsynchTextFrame) SetLanguage(language string) error {
	if len(language) != 3 {
		return errors.New("language: invalid language string")
	}

	f.language = language
	f.changeSize(0)
	return nil
}

func (f UnsynchTextFrame) String() string {
	return fmt.Sprintf("%s\t%s:\n%s", f.language, f.description, f.text)
}

func (f UnsynchTextFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteByte(f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.language, encodedbytes.NativeEncoding); err != nil {
		return bytes
	}

	if err = wr.WriteNullTermString(f.description, f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.text, f.encoding); err != nil {
		return bytes
	}

	return bytes
}

// ImageFrame represent frames that have media attached
type ImageFrame struct {
	DataFrame
	encoding    byte
	mimeType    string
	pictureType byte
	description string
}

func ParseImageFrame(head FrameHead, data []byte) Framer {
	var err error
	f := new(ImageFrame)
	f.FrameHead = head
	rd := encodedbytes.NewReader(data)

	if f.encoding, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.mimeType, err = rd.ReadNullTermString(encodedbytes.NativeEncoding); err != nil {
		return nil
	}

	if f.pictureType, err = rd.ReadByte(); err != nil {
		return nil
	}

	if f.description, err = rd.ReadNullTermString(f.encoding); err != nil {
		return nil
	}

	if f.data, err = rd.ReadRest(); err != nil {
		return nil
	}

	return f
}

func (f ImageFrame) Encoding() string {
	return encodedbytes.EncodingForIndex(f.encoding)
}

func (f *ImageFrame) SetEncoding(encoding string) error {
	i := byte(encodedbytes.IndexForEncoding(encoding))
	if i < 0 {
		return errors.New("encoding: invalid encoding")
	}

	diff, err := encodedbytes.EncodedDiff(i, f.description, f.encoding, f.description)
	if err != nil {
		return err
	}

	f.changeSize(diff)
	f.encoding = i
	return nil
}

func (f ImageFrame) MIMEType() string {
	return f.mimeType
}

func (f *ImageFrame) SetMIMEType(mimeType string) {
	diff := len(mimeType) - len(f.mimeType)
	if mimeType[len(mimeType)-1] != 0 {
		nullTermBytes := append([]byte(mimeType), 0x00)
		f.mimeType = string(nullTermBytes)
		diff += 1
	} else {
		f.mimeType = mimeType
	}

	f.changeSize(diff)
}

func (f ImageFrame) String() string {
	return fmt.Sprintf("%s\t%s: <binary data>", f.mimeType, f.description)
}

func (f ImageFrame) Bytes() []byte {
	var err error
	bytes := make([]byte, f.Size())
	wr := encodedbytes.NewWriter(bytes)

	if err = wr.WriteByte(f.encoding); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.mimeType, encodedbytes.NativeEncoding); err != nil {
		return bytes
	}

	if err = wr.WriteByte(f.pictureType); err != nil {
		return bytes
	}

	if err = wr.WriteString(f.description, f.encoding); err != nil {
		return bytes
	}

	if n, err := wr.Write(f.data); n < len(f.data) || err != nil {
		return bytes
	}

	return bytes
}
<|endoftext|>


<commit_before>package cruncy

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestFields(t *testing.T) {
	tmr := NewTimer("Test")
	tmr.Tick()
	fields := tmr.LogFields()
	assert.NotNil(t, fields)
	assert.Equal(t, "Test", fields["title"])
	assert.Equal(t, int64(1), fields["total_rows"])
}

func TestTimer(t *testing.T) {
	tmr := NewTimer("test2")
	tmr.Start()
	tmr.Tick()
	<-time.After(100 * time.Millisecond)
	tmr.Stop()
	tmr.ShowBatchTime()
	tmr.ShowTotalDuration()
	assert.Equal(t, int64(1), tmr.Index.Get())
	assert.Equal(t, int64(0), tmr.BatchDuractionSeconds())
	assert.Equal(t, int64(0), tmr.TotalDuractionSeconds())
	x := tmr.TotalDuration()
	assert.True(t, x.Seconds() < float64(1.0))
}
<commit_msg>Added simple benchmark for timer<commit_after>package cruncy

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestFields(t *testing.T) {
	tmr := NewTimer("Test")
	tmr.Tick()
	fields := tmr.LogFields()
	assert.NotNil(t, fields)
	assert.Equal(t, "Test", fields["title"])
	assert.Equal(t, int64(1), fields["total_rows"])
}

func TestTimer(t *testing.T) {
	tmr := NewTimer("test2")
	tmr.Start()
	tmr.Tick()
	<-time.After(100 * time.Millisecond)
	tmr.Stop()
	tmr.ShowBatchTime()
	tmr.ShowTotalDuration()
	assert.Equal(t, int64(1), tmr.Index.Get())
	assert.Equal(t, int64(0), tmr.BatchDuractionSeconds())
	assert.Equal(t, int64(0), tmr.TotalDuractionSeconds())
	x := tmr.TotalDuration()
	assert.True(t, x.Seconds() < float64(1.0))
}

func BenchmarkTimer(b *testing.B) {
	tmr := NewTimer("test2")
	tmr.Start()
	testRounds := int64(1000000)
	tmr.BatchSize = testRounds * int64(b.N)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		for i := int64(0); i < testRounds; i++ {
			tmr.Tick()
		}
	}
	tmr.Stop()

}
<|endoftext|>


<commit_before>package scheduler

import (
	"fmt"
	"strconv"

	memdb "github.com/hashicorp/go-memdb"
	"github.com/hashicorp/nomad/helper"
	"github.com/hashicorp/nomad/nomad/structs"
)

// propertySet is used to track the values used for a particular property.
type propertySet struct {
	// ctx is used to lookup the plan and state
	ctx Context

	// jobID is the job we are operating on
	jobID string

	// namespace is the namespace of the job we are operating on
	namespace string

	// taskGroup is optionally set if the constraint is for a task group
	taskGroup string

	// constraint is the constraint this property set is checking
	constraint *structs.Constraint

	// allowedCount is the allowed number of allocations that can have the
	// distinct property
	allowedCount uint64

	// errorBuilding marks whether there was an error when building the property
	// set
	errorBuilding error

	// existingValues is a mapping of the values of a property to the number of
	// times the value has been used by pre-existing allocations.
	existingValues map[string]uint64

	// proposedValues is a mapping of the values of a property to the number of
	// times the value has been used by proposed allocations.
	proposedValues map[string]uint64

	// clearedValues is a mapping of the values of a property to the number of
	// times the value has been used by proposed stopped allocations.
	clearedValues map[string]uint64
}

// NewPropertySet returns a new property set used to guarantee unique property
// values for new allocation placements.
func NewPropertySet(ctx Context, job *structs.Job) *propertySet {
	p := &propertySet{
		ctx:            ctx,
		jobID:          job.ID,
		namespace:      job.Namespace,
		existingValues: make(map[string]uint64),
	}

	return p
}

// SetJobConstraint is used to parameterize the property set for a
// distinct_property constraint set at the job level.
func (p *propertySet) SetJobConstraint(constraint *structs.Constraint) {
	p.setConstraint(constraint, "")
}

// SetTGConstraint is used to parameterize the property set for a
// distinct_property constraint set at the task group level. The inputs are the
// constraint and the task group name.
func (p *propertySet) SetTGConstraint(constraint *structs.Constraint, taskGroup string) {
	p.setConstraint(constraint, taskGroup)
}

// setConstraint is a shared helper for setting a job or task group constraint.
func (p *propertySet) setConstraint(constraint *structs.Constraint, taskGroup string) {
	// Store that this is for a task group
	if taskGroup != "" {
		p.taskGroup = taskGroup
	}

	// Store the constraint
	p.constraint = constraint

	// Determine the number of allowed allocations with the property.
	if v := constraint.RTarget; v != "" {
		c, err := strconv.ParseUint(v, 10, 64)
		if err != nil {
			p.errorBuilding = fmt.Errorf("failed to convert RTarget %q to uint64: %v", v, err)
			p.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: %v", p.errorBuilding)
			return
		}

		p.allowedCount = c
	} else {
		p.allowedCount = 1
	}

	// Determine the number of existing allocations that are using a property
	// value
	p.populateExisting(constraint)

	// Populate the proposed when setting the constraint. We do this because
	// when detecting if we can inplace update an allocation we stage an
	// eviction and then select. This means the plan has an eviction before a
	// single select has finished.
	p.PopulateProposed()
}

// populateExisting is a helper shared when setting the constraint to populate
// the existing values.
func (p *propertySet) populateExisting(constraint *structs.Constraint) {
	// Retrieve all previously placed allocations
	ws := memdb.NewWatchSet()
	allocs, err := p.ctx.State().AllocsByJob(ws, p.namespace, p.jobID, false)
	if err != nil {
		p.errorBuilding = fmt.Errorf("failed to get job's allocations: %v", err)
		p.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: %v", p.errorBuilding)
		return
	}

	// Filter to the correct set of allocs
	allocs = filterAllocs(allocs, true, p.taskGroup)

	// Get all the nodes that have been used by the allocs
	nodes, err := buildNodeMap(p.ctx.State(), allocs)
	if err != nil {
		p.errorBuilding = err
		p.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: %v", err)
		return
	}

	// Build existing properties map
	populateProperties(p.constraint.LTarget, allocs, nodes, p.existingValues)
}

// PopulateProposed populates the proposed values and recomputes any cleared
// value. It should be called whenever the plan is updated to ensure correct
// results when checking an option.
func (p *propertySet) PopulateProposed() {

	// Reset the proposed properties
	p.proposedValues = make(map[string]uint64)
	p.clearedValues = make(map[string]uint64)

	// Gather the set of proposed stops.
	var stopping []*structs.Allocation
	for _, updates := range p.ctx.Plan().NodeUpdate {
		stopping = append(stopping, updates...)
	}
	stopping = filterAllocs(stopping, false, p.taskGroup)

	// Gather the proposed allocations
	var proposed []*structs.Allocation
	for _, pallocs := range p.ctx.Plan().NodeAllocation {
		proposed = append(proposed, pallocs...)
	}
	proposed = filterAllocs(proposed, true, p.taskGroup)

	// Get the used nodes
	both := make([]*structs.Allocation, 0, len(stopping)+len(proposed))
	both = append(both, stopping...)
	both = append(both, proposed...)
	nodes, err := buildNodeMap(p.ctx.State(), both)
	if err != nil {
		p.errorBuilding = err
		p.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: %v", err)
		return
	}

	// Populate the cleared values
	populateProperties(p.constraint.LTarget, stopping, nodes, p.clearedValues)

	// Populate the proposed values
	populateProperties(p.constraint.LTarget, proposed, nodes, p.proposedValues)

	// Remove any cleared value that is now being used by the proposed allocs
	for value := range p.proposedValues {
		current, ok := p.clearedValues[value]
		if !ok {
			continue
		} else if current == 0 {
			delete(p.clearedValues, value)
		} else if current > 1 {
			p.clearedValues[value]--
		}
	}
}

// SatisfiesDistinctProperties checks if the option satisfies the
// distinct_property constraints given the existing placements and proposed
// placements. If the option does not satisfy the constraints an explanation is
// given.
func (p *propertySet) SatisfiesDistinctProperties(option *structs.Node, tg string) (bool, string) {
	// Check if there was an error building
	if p.errorBuilding != nil {
		return false, p.errorBuilding.Error()
	}

	// Get the nodes property value
	nValue, ok := getProperty(option, p.constraint.LTarget)
	if !ok {
		return false, fmt.Sprintf("missing property %q", p.constraint.LTarget)
	}

	// combine the counts of how many times the property has been used by
	// existing and proposed allocations
	combinedUse := make(map[string]uint64, helper.IntMax(len(p.existingValues), len(p.proposedValues)))
	for _, usedValues := range []map[string]uint64{p.existingValues, p.proposedValues} {
		for propertyValue, usedCount := range usedValues {
			combinedUse[propertyValue] += usedCount
		}
	}

	// Go through and discount the combined count when the value has been
	// cleared by a proposed stop.
	for propertyValue, clearedCount := range p.clearedValues {
		combined, ok := combinedUse[propertyValue]
		if !ok {
			continue
		}

		// Don't clear below 0.
		if combined >= clearedCount {
			combinedUse[propertyValue] = combined - clearedCount
		} else {
			combinedUse[propertyValue] = 0
		}
	}

	usedCount, used := combinedUse[nValue]
	if !used {
		// The property value has never been used so we can use it.
		return true, ""
	}

	// The property value has been used but within the number of allowed
	// allocations.
	if usedCount < p.allowedCount {
		return true, ""
	}

	return false, fmt.Sprintf("distinct_property: %s=%s used by %d allocs", p.constraint.LTarget, nValue, usedCount)
}

// filterAllocs filters a set of allocations to just be those that are running
// and if the property set is operation at a task group level, for allocations
// for that task group
func filterAllocs(allocs []*structs.Allocation, filterTerminal bool, taskGroup string) []*structs.Allocation {
	n := len(allocs)
	for i := 0; i < n; i++ {
		remove := false
		if filterTerminal {
			remove = allocs[i].TerminalStatus()
		}

		// If the constraint is on the task group filter the allocations to just
		// those on the task group
		if taskGroup != "" {
			remove = remove || allocs[i].TaskGroup != taskGroup
		}

		if remove {
			allocs[i], allocs[n-1] = allocs[n-1], nil
			i--
			n--
		}
	}
	return allocs[:n]
}

// buildNodeMap takes a list of allocations and returns a map of the nodes used
// by those allocations
func buildNodeMap(state State, allocs []*structs.Allocation) (map[string]*structs.Node, error) {
	// Get all the nodes that have been used by the allocs
	nodes := make(map[string]*structs.Node)
	ws := memdb.NewWatchSet()
	for _, alloc := range allocs {
		if _, ok := nodes[alloc.NodeID]; ok {
			continue
		}

		node, err := state.NodeByID(ws, alloc.NodeID)
		if err != nil {
			return nil, fmt.Errorf("failed to lookup node ID %q: %v", alloc.NodeID, err)
		}

		nodes[alloc.NodeID] = node
	}

	return nodes, nil
}

// populateProperties goes through all allocations and builds up the used
// properties from the nodes storing the results in the passed properties map.
func populateProperties(lTarget string, allocs []*structs.Allocation, nodes map[string]*structs.Node,
	properties map[string]uint64) {

	for _, alloc := range allocs {
		nProperty, ok := getProperty(nodes[alloc.NodeID], lTarget)
		if !ok {
			continue
		}

		properties[nProperty]++
	}
}

// getProperty is used to lookup the property value on the node
func getProperty(n *structs.Node, property string) (string, bool) {
	if n == nil || property == "" {
		return "", false
	}

	val, ok := resolveTarget(property, n)
	if !ok {
		return "", false
	}
	nodeValue, ok := val.(string)
	if !ok {
		return "", false
	}

	return nodeValue, true
}
<commit_msg>Back out changes to propertyset that were not necessary for affinities<commit_after>package scheduler

import (
	"fmt"
	"strconv"

	memdb "github.com/hashicorp/go-memdb"
	"github.com/hashicorp/nomad/helper"
	"github.com/hashicorp/nomad/nomad/structs"
)

// propertySet is used to track the values used for a particular property.
type propertySet struct {
	// ctx is used to lookup the plan and state
	ctx Context

	// jobID is the job we are operating on
	jobID string

	// namespace is the namespace of the job we are operating on
	namespace string

	// taskGroup is optionally set if the constraint is for a task group
	taskGroup string

	// constraint is the constraint this property set is checking
	constraint *structs.Constraint

	// allowedCount is the allowed number of allocations that can have the
	// distinct property
	allowedCount uint64

	// errorBuilding marks whether there was an error when building the property
	// set
	errorBuilding error

	// existingValues is a mapping of the values of a property to the number of
	// times the value has been used by pre-existing allocations.
	existingValues map[string]uint64

	// proposedValues is a mapping of the values of a property to the number of
	// times the value has been used by proposed allocations.
	proposedValues map[string]uint64

	// clearedValues is a mapping of the values of a property to the number of
	// times the value has been used by proposed stopped allocations.
	clearedValues map[string]uint64
}

// NewPropertySet returns a new property set used to guarantee unique property
// values for new allocation placements.
func NewPropertySet(ctx Context, job *structs.Job) *propertySet {
	p := &propertySet{
		ctx:            ctx,
		jobID:          job.ID,
		namespace:      job.Namespace,
		existingValues: make(map[string]uint64),
	}

	return p
}

// SetJobConstraint is used to parameterize the property set for a
// distinct_property constraint set at the job level.
func (p *propertySet) SetJobConstraint(constraint *structs.Constraint) {
	p.setConstraint(constraint, "")
}

// SetTGConstraint is used to parameterize the property set for a
// distinct_property constraint set at the task group level. The inputs are the
// constraint and the task group name.
func (p *propertySet) SetTGConstraint(constraint *structs.Constraint, taskGroup string) {
	p.setConstraint(constraint, taskGroup)
}

// setConstraint is a shared helper for setting a job or task group constraint.
func (p *propertySet) setConstraint(constraint *structs.Constraint, taskGroup string) {
	// Store that this is for a task group
	if taskGroup != "" {
		p.taskGroup = taskGroup
	}

	// Store the constraint
	p.constraint = constraint

	// Determine the number of allowed allocations with the property.
	if v := constraint.RTarget; v != "" {
		c, err := strconv.ParseUint(v, 10, 64)
		if err != nil {
			p.errorBuilding = fmt.Errorf("failed to convert RTarget %q to uint64: %v", v, err)
			p.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: %v", p.errorBuilding)
			return
		}

		p.allowedCount = c
	} else {
		p.allowedCount = 1
	}

	// Determine the number of existing allocations that are using a property
	// value
	p.populateExisting(constraint)

	// Populate the proposed when setting the constraint. We do this because
	// when detecting if we can inplace update an allocation we stage an
	// eviction and then select. This means the plan has an eviction before a
	// single select has finished.
	p.PopulateProposed()
}

// populateExisting is a helper shared when setting the constraint to populate
// the existing values.
func (p *propertySet) populateExisting(constraint *structs.Constraint) {
	// Retrieve all previously placed allocations
	ws := memdb.NewWatchSet()
	allocs, err := p.ctx.State().AllocsByJob(ws, p.namespace, p.jobID, false)
	if err != nil {
		p.errorBuilding = fmt.Errorf("failed to get job's allocations: %v", err)
		p.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: %v", p.errorBuilding)
		return
	}

	// Filter to the correct set of allocs
	allocs = p.filterAllocs(allocs, true)

	// Get all the nodes that have been used by the allocs
	nodes, err := p.buildNodeMap(allocs)
	if err != nil {
		p.errorBuilding = err
		p.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: %v", err)
		return
	}

	// Build existing properties map
	p.populateProperties(allocs, nodes, p.existingValues)
}

// PopulateProposed populates the proposed values and recomputes any cleared
// value. It should be called whenever the plan is updated to ensure correct
// results when checking an option.
func (p *propertySet) PopulateProposed() {

	// Reset the proposed properties
	p.proposedValues = make(map[string]uint64)
	p.clearedValues = make(map[string]uint64)

	// Gather the set of proposed stops.
	var stopping []*structs.Allocation
	for _, updates := range p.ctx.Plan().NodeUpdate {
		stopping = append(stopping, updates...)
	}
	stopping = p.filterAllocs(stopping, false)

	// Gather the proposed allocations
	var proposed []*structs.Allocation
	for _, pallocs := range p.ctx.Plan().NodeAllocation {
		proposed = append(proposed, pallocs...)
	}
	proposed = p.filterAllocs(proposed, true)

	// Get the used nodes
	both := make([]*structs.Allocation, 0, len(stopping)+len(proposed))
	both = append(both, stopping...)
	both = append(both, proposed...)
	nodes, err := p.buildNodeMap(both)
	if err != nil {
		p.errorBuilding = err
		p.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: %v", err)
		return
	}

	// Populate the cleared values
	p.populateProperties(stopping, nodes, p.clearedValues)

	// Populate the proposed values
	p.populateProperties(proposed, nodes, p.proposedValues)

	// Remove any cleared value that is now being used by the proposed allocs
	for value := range p.proposedValues {
		current, ok := p.clearedValues[value]
		if !ok {
			continue
		} else if current == 0 {
			delete(p.clearedValues, value)
		} else if current > 1 {
			p.clearedValues[value]--
		}
	}
}

// SatisfiesDistinctProperties checks if the option satisfies the
// distinct_property constraints given the existing placements and proposed
// placements. If the option does not satisfy the constraints an explanation is
// given.
func (p *propertySet) SatisfiesDistinctProperties(option *structs.Node, tg string) (bool, string) {
	// Check if there was an error building
	if p.errorBuilding != nil {
		return false, p.errorBuilding.Error()
	}

	// Get the nodes property value
	nValue, ok := getProperty(option, p.constraint.LTarget)
	if !ok {
		return false, fmt.Sprintf("missing property %q", p.constraint.LTarget)
	}

	// combine the counts of how many times the property has been used by
	// existing and proposed allocations
	combinedUse := make(map[string]uint64, helper.IntMax(len(p.existingValues), len(p.proposedValues)))
	for _, usedValues := range []map[string]uint64{p.existingValues, p.proposedValues} {
		for propertyValue, usedCount := range usedValues {
			combinedUse[propertyValue] += usedCount
		}
	}

	// Go through and discount the combined count when the value has been
	// cleared by a proposed stop.
	for propertyValue, clearedCount := range p.clearedValues {
		combined, ok := combinedUse[propertyValue]
		if !ok {
			continue
		}

		// Don't clear below 0.
		if combined >= clearedCount {
			combinedUse[propertyValue] = combined - clearedCount
		} else {
			combinedUse[propertyValue] = 0
		}
	}

	usedCount, used := combinedUse[nValue]
	if !used {
		// The property value has never been used so we can use it.
		return true, ""
	}

	// The property value has been used but within the number of allowed
	// allocations.
	if usedCount < p.allowedCount {
		return true, ""
	}

	return false, fmt.Sprintf("distinct_property: %s=%s used by %d allocs", p.constraint.LTarget, nValue, usedCount)
}

// filterAllocs filters a set of allocations to just be those that are running
// and if the property set is operation at a task group level, for allocations
// for that task group
func (p *propertySet) filterAllocs(allocs []*structs.Allocation, filterTerminal bool) []*structs.Allocation {
	n := len(allocs)
	for i := 0; i < n; i++ {
		remove := false
		if filterTerminal {
			remove = allocs[i].TerminalStatus()
		}

		// If the constraint is on the task group filter the allocations to just
		// those on the task group
		if p.taskGroup != "" {
			remove = remove || allocs[i].TaskGroup != p.taskGroup
		}

		if remove {
			allocs[i], allocs[n-1] = allocs[n-1], nil
			i--
			n--
		}
	}
	return allocs[:n]
}

// buildNodeMap takes a list of allocations and returns a map of the nodes used
// by those allocations
func (p *propertySet) buildNodeMap(allocs []*structs.Allocation) (map[string]*structs.Node, error) {
	// Get all the nodes that have been used by the allocs
	nodes := make(map[string]*structs.Node)
	ws := memdb.NewWatchSet()
	for _, alloc := range allocs {
		if _, ok := nodes[alloc.NodeID]; ok {
			continue
		}

		node, err := p.ctx.State().NodeByID(ws, alloc.NodeID)
		if err != nil {
			return nil, fmt.Errorf("failed to lookup node ID %q: %v", alloc.NodeID, err)
		}

		nodes[alloc.NodeID] = node
	}

	return nodes, nil
}

// populateProperties goes through all allocations and builds up the used
// properties from the nodes storing the results in the passed properties map.
func (p *propertySet) populateProperties(allocs []*structs.Allocation, nodes map[string]*structs.Node,
	properties map[string]uint64) {

	for _, alloc := range allocs {
		nProperty, ok := getProperty(nodes[alloc.NodeID], p.constraint.LTarget)
		if !ok {
			continue
		}

		properties[nProperty]++
	}
}

// getProperty is used to lookup the property value on the node
func getProperty(n *structs.Node, property string) (string, bool) {
	if n == nil || property == "" {
		return "", false
	}

	val, ok := resolveTarget(property, n)
	if !ok {
		return "", false
	}
	nodeValue, ok := val.(string)
	if !ok {
		return "", false
	}

	return nodeValue, true
}
<|endoftext|>


<commit_before>package feed_test

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"github.com/Forau/yanngo/feed"

	"math/big"
	"net"
	"testing"
	"time"
)

// We generate a self signed cert, just for this test run.
var TLS *tls.Config

func init() {
	ca := &x509.Certificate{
		SerialNumber: big.NewInt(1812),
		Subject: pkix.Name{
			Country:            []string{"Internet"},
			Organization:       []string{"The Internet"},
			OrganizationalUnit: []string{"Moving bits"},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(10, 0, 0),
		SubjectKeyId:          []byte{8, 8, 8, 8, 8},
		BasicConstraintsValid: true,
		IsCA:        true,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		KeyUsage:    x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		IPAddresses: []net.IP{net.ParseIP("127.0.0.1")},
	}
	priv, _ := rsa.GenerateKey(rand.Reader, 1024)
	pub := &priv.PublicKey
	ca_b, err := x509.CreateCertificate(rand.Reader, ca, ca, pub, priv)
	if err != nil {
		panic("create ca failed: " + err.Error())
	}

	//  certPair,_ := tls.X509KeyPair(ca_b,x509.MarshalPKCS1PrivateKey(priv))
	certPair := tls.Certificate{
		Certificate: [][]byte{ca_b},
		PrivateKey:  priv,
	}

	pool := x509.NewCertPool()
	pool.AddCert(ca)

	// This is the server cert. Self signed, and not very valid, so client MUST use InsecureSkipVerify
	TLS = &tls.Config{
		Certificates: []tls.Certificate{certPair},
		RootCAs:      pool,
	}

	// To make the feed accept self signed certs
	feed.DefaultTLS = &tls.Config{InsecureSkipVerify: true}
}

type testSrv struct {
	listen   net.Listener
	t        *testing.T
	exit     chan interface{}
	isClosed bool

	connFn func(c net.Conn)
}

func (ts *testSrv) Close() error {
	close(ts.exit)
	for i := 0; i < 10; i++ {
		if ts.isClosed {
			return nil
		}
		time.Sleep(100 * time.Millisecond)
	}
	return fmt.Errorf("Server did not close within one second...")
}

func (ts *testSrv) mainLoop() {
	defer ts.listen.Close()
	go func(l net.Listener) {
		for {
			conn, err := l.Accept()
			if err != nil {
				ts.t.Log(err)
				return
			}
			go ts.connFn(conn)
		}
	}(ts.listen)
	<-ts.exit // Wait for exit
	ts.isClosed = true
	ts.t.Logf("Exiting main loop for %+v", ts)
}

func newTestSrv(t *testing.T, connFn func(net.Conn)) (srv *testSrv) {
	// &tls.Config{RootCAs: TLS.RootCAs}
	l, err := tls.Listen("tcp", "127.0.0.1:", TLS)
	if err != nil {
		t.Fatal(err)
	}
	srv = &testSrv{
		listen: l,
		t:      t,
		exit:   make(chan interface{}),
		connFn: connFn,
	}
	go srv.mainLoop()
	t.Logf("New test server: %+v, listening on %s", srv, srv.listen.Addr().String())
	return
}

type simpleCallback struct {
	t *testing.T
}

func (c *simpleCallback) OnConnect(w feed.CmdWriter, ft feed.FeedType) {
	c.t.Logf("Connect[%v]: %+v", ft, w)
}
func (c *simpleCallback) OnMessage(msg *feed.FeedMsg, ft feed.FeedType) {
	c.t.Logf("Msg[%v]: %+v", ft, msg.String())
}
func (c *simpleCallback) OnError(err error, ft feed.FeedType) { c.t.Logf("Err[%v]: %+v", ft, err) }

func TestConnectToFeed(t *testing.T) {
	quit := make(chan interface{})
	privts := newTestSrv(t, func(c net.Conn) {
		defer c.Close()
		t.Log("Got connection: ", c)
		for {
			buff := make([]byte, 1024)
			n, err := c.Read(buff)
			t.Log("Read %d bytes(%+v): %s", n, err, string(buff))
			if err != nil {
				return
			}
		}
	})
	defer privts.Close()
	privSess := func() (key, url string, err error) {
		return "PRIV", privts.listen.Addr().String(), nil
	}

	pubts := newTestSrv(t, func(c net.Conn) {
		defer c.Close()
		t.Log("Got connection: ", c)

		for {
			buff := make([]byte, 1024)
			n, err := c.Read(buff)
			t.Log("Read %d bytes(%+v): %s", n, err, string(buff))
			if err != nil {
				return
			}
		}
	})
	defer pubts.Close()
	pubSess := func() (key, url string, err error) {
		return "PUB", pubts.listen.Addr().String(), nil
	}

	cb := &simpleCallback{t}

	fd, err := feed.NewFeedDaemon(privSess, pubSess, cb)
	if err != nil {
		t.Fatalf("Daemon error: %+v", err)
	}

	time.Sleep(10 * time.Millisecond) // Ugly yield

	err = fd.Subscribe("price", "46", "11")
	if err != nil {
		t.Fatalf("Subscribe error: %+v", err)
	}
	fmt.Println("Closing: ", fd.Close())
	close(quit)
}
<commit_msg>Fixed nicer 'wait' for feed_test<commit_after>package feed_test

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"github.com/Forau/yanngo/feed"

	"math/big"
	"net"
	"testing"
	"time"
)

// We generate a self signed cert, just for this test run.
var TLS *tls.Config

func init() {
	ca := &x509.Certificate{
		SerialNumber: big.NewInt(1812),
		Subject: pkix.Name{
			Country:            []string{"Internet"},
			Organization:       []string{"The Internet"},
			OrganizationalUnit: []string{"Moving bits"},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(10, 0, 0),
		SubjectKeyId:          []byte{8, 8, 8, 8, 8},
		BasicConstraintsValid: true,
		IsCA:        true,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		KeyUsage:    x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		IPAddresses: []net.IP{net.ParseIP("127.0.0.1")},
	}
	priv, _ := rsa.GenerateKey(rand.Reader, 1024)
	pub := &priv.PublicKey
	ca_b, err := x509.CreateCertificate(rand.Reader, ca, ca, pub, priv)
	if err != nil {
		panic("create ca failed: " + err.Error())
	}

	//  certPair,_ := tls.X509KeyPair(ca_b,x509.MarshalPKCS1PrivateKey(priv))
	certPair := tls.Certificate{
		Certificate: [][]byte{ca_b},
		PrivateKey:  priv,
	}

	pool := x509.NewCertPool()
	pool.AddCert(ca)

	// This is the server cert. Self signed, and not very valid, so client MUST use InsecureSkipVerify
	TLS = &tls.Config{
		Certificates: []tls.Certificate{certPair},
		RootCAs:      pool,
	}

	// To make the feed accept self signed certs
	feed.DefaultTLS = &tls.Config{InsecureSkipVerify: true}
}

type testSrv struct {
	listen   net.Listener
	t        *testing.T
	exit     chan interface{}
	isClosed bool

	connFn func(c net.Conn)
}

func (ts *testSrv) Close() error {
	close(ts.exit)
	for i := 0; i < 10; i++ {
		if ts.isClosed {
			return nil
		}
		time.Sleep(100 * time.Millisecond)
	}
	return fmt.Errorf("Server did not close within one second...")
}

func (ts *testSrv) mainLoop() {
	defer ts.listen.Close()
	go func(l net.Listener) {
		for {
			conn, err := l.Accept()
			if err != nil {
				ts.t.Log(err)
				return
			}
			go ts.connFn(conn)
		}
	}(ts.listen)
	<-ts.exit // Wait for exit
	ts.isClosed = true
	ts.t.Logf("Exiting main loop for %+v", ts)
}

func newTestSrv(t *testing.T, connFn func(net.Conn)) (srv *testSrv) {
	// &tls.Config{RootCAs: TLS.RootCAs}
	l, err := tls.Listen("tcp", "127.0.0.1:0", TLS)
	if err != nil {
		t.Fatal(err)
	}
	srv = &testSrv{
		listen: l,
		t:      t,
		exit:   make(chan interface{}),
		connFn: connFn,
	}
	go srv.mainLoop()
	t.Logf("New test server: %+v, listening on %s", srv, srv.listen.Addr().String())
	return
}

type simpleCallback struct {
	t           *testing.T
	connectChan chan bool
}

func (c *simpleCallback) OnConnect(w feed.CmdWriter, ft feed.FeedType) {
	c.t.Logf("Connect[%v]: %+v", ft, w)
	c.connectChan <- true
}
func (c *simpleCallback) OnMessage(msg *feed.FeedMsg, ft feed.FeedType) {
	c.t.Logf("Msg[%v]: %+v", ft, msg.String())
}
func (c *simpleCallback) OnError(err error, ft feed.FeedType) { c.t.Logf("Err[%v]: %+v", ft, err) }

func TestConnectToFeed(t *testing.T) {
	quit := make(chan interface{})
	privts := newTestSrv(t, func(c net.Conn) {
		defer c.Close()
		t.Log("Got connection: ", c)
		for {
			buff := make([]byte, 1024)
			n, err := c.Read(buff)
			t.Log("Read %d bytes(%+v): %s", n, err, string(buff))
			if err != nil {
				return
			}
		}
	})
	defer privts.Close()
	privSess := func() (key, url string, err error) {
		return "PRIV", privts.listen.Addr().String(), nil
	}

	pubts := newTestSrv(t, func(c net.Conn) {
		defer c.Close()
		t.Log("Got connection: ", c)

		for {
			buff := make([]byte, 1024)
			n, err := c.Read(buff)
			t.Log("Read %d bytes(%+v): %s", n, err, string(buff))
			if err != nil {
				return
			}
		}
	})
	defer pubts.Close()
	pubSess := func() (key, url string, err error) {
		return "PUB", pubts.listen.Addr().String(), nil
	}

	cb := &simpleCallback{t, make(chan bool)}

	fd, err := feed.NewFeedDaemon(privSess, pubSess, cb)
	if err != nil {
		t.Fatalf("Daemon error: %+v", err)
	}

	go func() {
		<-cb.connectChan
		<-cb.connectChan

		err = fd.Subscribe("price", "46", "11")
		if err != nil {
			t.Fatalf("Subscribe error: %+v", err)
		}
		fmt.Println("Closing: ", fd.Close())
		close(quit)
	}()
	<-quit
}
<|endoftext|>


<commit_before>package signal

import (
	"context"
	"time"
)

type ActivityUpdater interface {
	Update()
}

type ActivityTimer struct {
	updated chan bool
	timeout chan time.Duration
	ctx     context.Context
	cancel  context.CancelFunc
}

func (t *ActivityTimer) Update() {
	select {
	case t.updated <- true:
	default:
	}
}

func (t *ActivityTimer) SetTimeout(timeout time.Duration) {
	t.timeout <- timeout
}

func (t *ActivityTimer) run() {
	ticker := time.NewTicker(<-t.timeout)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
		case <-t.ctx.Done():
			return
		case timeout := <-t.timeout:
			ticker.Stop()
			ticker = time.NewTicker(timeout)
		}

		select {
		case <-t.updated:
		// Updated keep waiting.
		default:
			t.cancel()
			return
		}
	}
}

func CancelAfterInactivity(ctx context.Context, timeout time.Duration) (context.Context, *ActivityTimer) {
	ctx, cancel := context.WithCancel(ctx)
	timer := &ActivityTimer{
		ctx:     ctx,
		cancel:  cancel,
		timeout: make(chan time.Duration, 1),
		updated: make(chan bool, 1),
	}
	timer.timeout <- timeout
	go timer.run()
	return ctx, timer
}
<commit_msg>fix timer leak<commit_after>package signal

import (
	"context"
	"time"
)

type ActivityUpdater interface {
	Update()
}

type ActivityTimer struct {
	updated chan bool
	timeout chan time.Duration
	ctx     context.Context
	cancel  context.CancelFunc
}

func (t *ActivityTimer) Update() {
	select {
	case t.updated <- true:
	default:
	}
}

func (t *ActivityTimer) SetTimeout(timeout time.Duration) {
	t.timeout <- timeout
}

func (t *ActivityTimer) run() {
	ticker := time.NewTicker(<-t.timeout)
	defer func() {
		ticker.Stop()
	}()

	for {
		select {
		case <-ticker.C:
		case <-t.ctx.Done():
			return
		case timeout := <-t.timeout:
			ticker.Stop()
			ticker = time.NewTicker(timeout)
		}

		select {
		case <-t.updated:
		// Updated keep waiting.
		default:
			t.cancel()
			return
		}
	}
}

func CancelAfterInactivity(ctx context.Context, timeout time.Duration) (context.Context, *ActivityTimer) {
	ctx, cancel := context.WithCancel(ctx)
	timer := &ActivityTimer{
		ctx:     ctx,
		cancel:  cancel,
		timeout: make(chan time.Duration, 1),
		updated: make(chan bool, 1),
	}
	timer.timeout <- timeout
	go timer.run()
	return ctx, timer
}
<|endoftext|>


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

// +build !windows

package cgotest

/*
#include <signal.h>
#include <pthread.h>
#include <unistd.h>
#include <stdlib.h>

static void *thread(void *p) {
	(void)p;
	const int M = 100;
	int i;
	for (i = 0; i < M; i++) {
		pthread_kill(pthread_self(), SIGCHLD);
		usleep(rand() % 20 + 5);
	}
	return NULL;
}
void testSendSIG() {
	const int N = 20;
	int i;
	pthread_t tid[N];
	for (i = 0; i < N; i++) {
		usleep(rand() % 200 + 100);
		pthread_create(&tid[i], 0, thread, NULL);
	}
	for (i = 0; i < N; i++)
		pthread_join(tid[i], 0);
}
*/
import "C"

import (
	"os"
	"os/signal"
	"syscall"
	"testing"
	"time"
)

func test3250(t *testing.T) {
	const (
		thres = 1
		sig   = syscall.SIGCHLD
	)
	type result struct {
		n   int
		sig os.Signal
	}
	var (
		sigCh     = make(chan os.Signal, 10)
		waitStart = make(chan struct{})
		waitDone  = make(chan result)
	)

	signal.Notify(sigCh, sig)

	go func() {
		n := 0
		alarm := time.After(time.Second * 3)
		for {
			select {
			case <-waitStart:
				waitStart = nil
			case v := <-sigCh:
				n++
				if v != sig || n > thres {
					waitDone <- result{n, v}
					return
				}
			case <-alarm:
				waitDone <- result{n, sig}
				return
			}
		}
	}()

	waitStart <- struct{}{}
	C.testSendSIG()
	r := <-waitDone
	if r.sig != sig {
		t.Fatalf("received signal %v, but want %v", r.sig, sig)
	}
	t.Logf("got %d signals\n", r.n)
	if r.n <= thres {
		t.Fatalf("expected more than %d", thres)
	}
}
<commit_msg>misc/cgo/test: skip test for issue 3250, the runtime isn't ready yet. see issue 5885.<commit_after>// Copyright 2013 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !windows

package cgotest

/*
#include <signal.h>
#include <pthread.h>
#include <unistd.h>
#include <stdlib.h>

static void *thread(void *p) {
	(void)p;
	const int M = 100;
	int i;
	for (i = 0; i < M; i++) {
		pthread_kill(pthread_self(), SIGCHLD);
		usleep(rand() % 20 + 5);
	}
	return NULL;
}
void testSendSIG() {
	const int N = 20;
	int i;
	pthread_t tid[N];
	for (i = 0; i < N; i++) {
		usleep(rand() % 200 + 100);
		pthread_create(&tid[i], 0, thread, NULL);
	}
	for (i = 0; i < N; i++)
		pthread_join(tid[i], 0);
}
*/
import "C"

import (
	"os"
	"os/signal"
	"syscall"
	"testing"
	"time"
)

func test3250(t *testing.T) {
	t.Skip("skipped, see golang.org/issue/5885")
	const (
		thres = 1
		sig   = syscall.SIGCHLD
	)
	type result struct {
		n   int
		sig os.Signal
	}
	var (
		sigCh     = make(chan os.Signal, 10)
		waitStart = make(chan struct{})
		waitDone  = make(chan result)
	)

	signal.Notify(sigCh, sig)

	go func() {
		n := 0
		alarm := time.After(time.Second * 3)
		for {
			select {
			case <-waitStart:
				waitStart = nil
			case v := <-sigCh:
				n++
				if v != sig || n > thres {
					waitDone <- result{n, v}
					return
				}
			case <-alarm:
				waitDone <- result{n, sig}
				return
			}
		}
	}()

	waitStart <- struct{}{}
	C.testSendSIG()
	r := <-waitDone
	if r.sig != sig {
		t.Fatalf("received signal %v, but want %v", r.sig, sig)
	}
	t.Logf("got %d signals\n", r.n)
	if r.n <= thres {
		t.Fatalf("expected more than %d", thres)
	}
}
<|endoftext|>


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

package str

import (
	"path/filepath"
	"strings"
)

// HasPath reports whether the slash-separated path s
// begins with the elements in prefix.
func HasPathPrefix(s, prefix string) bool {
	if len(s) == len(prefix) {
		return s == prefix
	}
	if prefix == "" {
		return true
	}
	if len(s) > len(prefix) {
		if prefix != "" && prefix[len(prefix)-1] == '/' || s[len(prefix)] == '/' {
			return s[:len(prefix)] == prefix
		}
	}
	return false
}

// HasFilePathPrefix reports whether the filesystem path s
// begins with the elements in prefix.
func HasFilePathPrefix(s, prefix string) bool {
	sv := strings.ToUpper(filepath.VolumeName(s))
	pv := strings.ToUpper(filepath.VolumeName(prefix))
	s = s[len(sv):]
	prefix = prefix[len(pv):]
	switch {
	default:
		return false
	case sv != pv:
		return false
	case len(s) == len(prefix):
		return s == prefix
	case prefix == "":
		return true
	case len(s) > len(prefix):
		if prefix[len(prefix)-1] == filepath.Separator {
			return strings.HasPrefix(s, prefix)
		}
		return s[len(prefix)] == filepath.Separator && s[:len(prefix)] == prefix
	}
}
<commit_msg>cmd/go/internal/str: simplify HasPathPrefix by epsilon<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package str

import (
	"path/filepath"
	"strings"
)

// HasPath reports whether the slash-separated path s
// begins with the elements in prefix.
func HasPathPrefix(s, prefix string) bool {
	if len(s) == len(prefix) {
		return s == prefix
	}
	if prefix == "" {
		return true
	}
	if len(s) > len(prefix) {
		if prefix[len(prefix)-1] == '/' || s[len(prefix)] == '/' {
			return s[:len(prefix)] == prefix
		}
	}
	return false
}

// HasFilePathPrefix reports whether the filesystem path s
// begins with the elements in prefix.
func HasFilePathPrefix(s, prefix string) bool {
	sv := strings.ToUpper(filepath.VolumeName(s))
	pv := strings.ToUpper(filepath.VolumeName(prefix))
	s = s[len(sv):]
	prefix = prefix[len(pv):]
	switch {
	default:
		return false
	case sv != pv:
		return false
	case len(s) == len(prefix):
		return s == prefix
	case prefix == "":
		return true
	case len(s) > len(prefix):
		if prefix[len(prefix)-1] == filepath.Separator {
			return strings.HasPrefix(s, prefix)
		}
		return s[len(prefix)] == filepath.Separator && s[:len(prefix)] == prefix
	}
}
<|endoftext|>


<commit_before>/**
 * @file scoreboard.go
 * @author Mikhail Klementyev jollheef<AT>riseup.net
 * @license GNU AGPLv3
 * @date November, 2015
 * @brief web scoreboard
 *
 * Contain web ui and several helpers for convert results to table
 */

package scoreboard

import (
	"database/sql"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"time"

	"github.com/jollheef/henhouse/game"
	"golang.org/x/net/websocket"
)

const (
	contestStateNotAvailable = "state n/a"
	contestNotStarted        = "not started"
	contestRunning           = "running"
	contestCompleted         = "completed"
)

var (
	gameShim      *game.Game
	contestStatus string
)

var (
	// InfoTimeout timeout between update info through websocket
	InfoTimeout = time.Second
	// ScoreboardTimeout timeout between update scoreboard through websocket
	ScoreboardTimeout = time.Second
	// TasksTimeout timeout between update tasks through websocket
	TasksTimeout = time.Second
	// FlagTimeout timeout between send flags
	FlagTimeout = time.Second
	// ScoreboardRecalcTimeout timeout between update scoreboard
	ScoreboardRecalcTimeout = time.Second
)

func durationToHMS(d time.Duration) string {

	sec := int(d.Seconds())

	var h, m, s int

	h = sec / 60 / 60
	m = (sec / 60) % 60
	s = sec % 60

	return fmt.Sprintf("%02d:%02d:%02d", h, m, s)
}

func getInfo() string {

	var left time.Duration
	var btnType string

	now := time.Now()

	if now.Before(gameShim.Start) {

		contestStatus = contestNotStarted
		left = gameShim.Start.Sub(now)
		btnType = "stop"

	} else if now.Before(gameShim.End) {

		contestStatus = contestRunning
		left = gameShim.End.Sub(now)
		btnType = "run"

	} else {
		contestStatus = contestCompleted
		left = 0
		btnType = "stop"
	}

	info := fmt.Sprintf(`<span id="game_status-%s">contest %s</span>`,
		btnType, contestStatus)

	if left != 0 {
		info += fmt.Sprintf(`<span id="timer">Left %s</span>`,
			durationToHMS(left))
	}

	return info
}

func infoHandler(ws *websocket.Conn) {

	defer ws.Close()
	for {
		_, err := fmt.Fprint(ws, getInfo())
		if err != nil {
			//log.Println("Socket closed:", err)
			return
		}

		time.Sleep(InfoTimeout)
	}
}

func scoreboardHandler(ws *websocket.Conn) {

	defer ws.Close()

	teamID := getTeamID(ws.Request())

	currentResult := scoreboardHTML(teamID)

	fmt.Fprint(ws, currentResult)

	sendedResult := currentResult

	lastUpdate := time.Now()

	for {
		currentResult := scoreboardHTML(teamID)

		if sendedResult != currentResult ||
			time.Now().After(lastUpdate.Add(time.Minute)) {

			sendedResult = currentResult
			lastUpdate = time.Now()

			_, err := fmt.Fprint(ws, currentResult)
			if err != nil {
				//log.Println("Socket closed:", err)
				return
			}
		}

		time.Sleep(ScoreboardTimeout)
	}
}

func scoreboardHTML(teamID int) (result string) {

	result = "
" +
		"#	" + "Team	" +
		"Score	" +
		"
"

	result += "
"

	scores, err := gameShim.Scoreboard()
	if err != nil {
		log.Println("Get scoreboard fail:", err)
		return
	}

	for n, teamScore := range scores {
		if teamScore.ID == teamID {
			result += `<tr class="self-team">`
		} else {
			result += `<tr>`
		}

		result += fmt.Sprintf(
			`<td class="team_index">%d</td>`+
				`<td class="team_name">%s</td>`+
				`<td class="team_score">%d</td></tr>`,
			n+1, teamScore.Name, teamScore.Score)

	}

	result += "
"

	return
}

func scoreboardUpdater(game *game.Game, updateTimeout time.Duration) {

	for {
		time.Sleep(updateTimeout)

		err := game.RecalcScoreboard()
		if err != nil {
			log.Println("Recalc scoreboard fail:", err)
		}
	}
}

func tasksHTML(teamID int) (result string) {

	cats, err := gameShim.Tasks()
	if err != nil {
		log.Println("Get tasks fail:", err)
	}

	for _, cat := range cats {
		result += categoryToHTML(teamID, cat)
	}

	return
}

func tasksHandler(ws *websocket.Conn) {

	defer ws.Close()

	teamID := getTeamID(ws.Request())

	currentTasks := tasksHTML(teamID)

	fmt.Fprint(ws, currentTasks)

	sendedTasks := currentTasks

	lastUpdate := time.Now()

	for {
		currentTasks := tasksHTML(teamID)

		if sendedTasks != currentTasks ||
			time.Now().After(lastUpdate.Add(time.Minute)) {

			sendedTasks = currentTasks
			lastUpdate = time.Now()

			_, err := fmt.Fprint(ws, currentTasks)
			if err != nil {
				//log.Println("Socket closed:", err)
				return
			}
		}

		time.Sleep(TasksTimeout)
	}
}

func taskHandler(w http.ResponseWriter, r *http.Request) {

	id, err := strconv.Atoi(r.URL.Query().Get("id"))
	if err != nil {
		log.Println("Atoi fail:", err)
		http.Redirect(w, r, "/", 307)
		return
	}

	cats, err := gameShim.Tasks()
	if err != nil {
		log.Println("Get tasks fail:", err)
		http.Redirect(w, r, "/", 307)
		return
	}

	task := game.TaskInfo{ID: id, Opened: false}

	for _, c := range cats {
		for _, t := range c.TasksInfo {
			if t.ID == id {
				task = t
				break
			}
		}
	}

	if !task.Opened {
		// Try to see closed task -> gtfo
		http.Redirect(w, r, "/", 307)
		return
	}

	teamID := getTeamID(r)

	flagSubmitFormat := `<br>` +
		`<form class="input-group" action="/flag?id=%d" method="post">` +
		`<input class="form-control float-left" name="flag" value="" placeholder="Flag">` +
		`<span class="input-group-btn">` +
		`<button class="btn btn-submit">Submit</button>` +
		`</span>` +
		`</form>`

	var submitForm string
	if taskSolvedBy(task, teamID) {
		submitForm = "Already solved"
	} else {
		submitForm = fmt.Sprintf(flagSubmitFormat, task.ID)
	}

	fmt.Fprintf(w, `<!DOCTYPE html>

    <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">

    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <link rel="shortcut icon" href="images/favicon.png" type="image/png">
    <title>Juniors CTF</title>

    <link rel="stylesheet" href="css/style.css" class="--apng-checked">

Scoreboard

Tasks

News

Sponsors

%s

        %s

        %s

          %s

`, task.Name, task.Desc, task.Author, submitForm)
}

func flagHandler(w http.ResponseWriter, r *http.Request) {

	if r.Method != "POST" {
		http.Redirect(w, r, "/", 307)
		return
	}

	taskID, err := strconv.Atoi(r.URL.Query().Get("id"))
	if err != nil {
		log.Println("Atoi fail:", err)
		http.Redirect(w, r, "/", 307)
		return
	}

	flag := r.FormValue("flag")

	teamID := getTeamID(r)

	solved, err := gameShim.Solve(teamID, taskID, flag)
	if err != nil {
		solved = false
	}

	var solvedMsg string
	if solved {
		solvedMsg = `<div class="flag_status solved">Solved</div>`
	} else {
		solvedMsg = `<div class="flag_status invalid">Invalid flag</div>`
	}

	log.Printf("Team ID: %d, Task ID: %d, Flag: %s, Result: %s\n",
		teamID, taskID, flag, solvedMsg)

	time.Sleep(FlagTimeout)

	fmt.Fprintf(w, `<!DOCTYPE html>

    <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">

    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <link rel="shortcut icon" href="images/favicon.png" type="image/png">
    <title>Juniors CTF</title>

    <link rel="stylesheet" href="css/style.css" class="--apng-checked">

Scoreboard

Tasks

News

Sponsors

%s

`, solvedMsg)
}

func handleStaticFile(pattern, file string) {
	http.HandleFunc(pattern,
		func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, file)
		})
}

func handleStaticFileSimple(file, wwwPath string) {
	handleStaticFile(file, wwwPath+file)
}

// Scoreboard implements web scoreboard
func Scoreboard(database *sql.DB, game *game.Game,
	wwwPath, templatePath, addr string) (err error) {

	contestStatus = contestStateNotAvailable
	gameShim = game
	templatePath = templatePath

	go scoreboardUpdater(game, ScoreboardRecalcTimeout)

	// Static files
	handleStaticFileSimple("/css/style.css", wwwPath)
	handleStaticFileSimple("/js/scoreboard.js", wwwPath)
	handleStaticFileSimple("/js/tasks.js", wwwPath)
	handleStaticFileSimple("/news.html", wwwPath)
	handleStaticFileSimple("/sponsors.html", wwwPath)
	handleStaticFileSimple("/images/bg.jpg", wwwPath)
	handleStaticFileSimple("/images/favicon.ico", wwwPath)
	handleStaticFileSimple("/images/favicon.png", wwwPath)
	handleStaticFileSimple("/images/401.jpg", wwwPath)
	handleStaticFileSimple("/images/juniors_ctf_txt.png", wwwPath)
	handleStaticFileSimple("/auth.html", wwwPath)

	// Get
	http.Handle("/", authorized(database, http.HandlerFunc(staticScoreboard)))
	http.Handle("/index.html", authorized(database, http.HandlerFunc(staticScoreboard)))
	http.Handle("/tasks.html", authorized(database, http.HandlerFunc(staticTasks)))
	http.Handle("/logout", authorized(database, http.HandlerFunc(logoutHandler)))

	// Websocket
	http.Handle("/scoreboard", authorized(database, websocket.Handler(scoreboardHandler)))
	http.Handle("/info", authorized(database, websocket.Handler(infoHandler)))
	http.Handle("/tasks", authorized(database, websocket.Handler(tasksHandler)))

	// Post
	http.Handle("/task", authorized(database, http.HandlerFunc(taskHandler)))
	http.Handle("/flag", authorized(database, http.HandlerFunc(flagHandler)))

	http.HandleFunc("/auth.php", http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			authHandler(database, w, r)
		}))

	log.Println("Launching scoreboard at", addr)

	return http.ListenAndServe(addr, nil)
}
<commit_msg>Fix set template path<commit_after>/**
 * @file scoreboard.go
 * @author Mikhail Klementyev jollheef<AT>riseup.net
 * @license GNU AGPLv3
 * @date November, 2015
 * @brief web scoreboard
 *
 * Contain web ui and several helpers for convert results to table
 */

package scoreboard

import (
	"database/sql"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"time"

	"github.com/jollheef/henhouse/game"
	"golang.org/x/net/websocket"
)

const (
	contestStateNotAvailable = "state n/a"
	contestNotStarted        = "not started"
	contestRunning           = "running"
	contestCompleted         = "completed"
)

var (
	gameShim      *game.Game
	contestStatus string
)

var (
	// InfoTimeout timeout between update info through websocket
	InfoTimeout = time.Second
	// ScoreboardTimeout timeout between update scoreboard through websocket
	ScoreboardTimeout = time.Second
	// TasksTimeout timeout between update tasks through websocket
	TasksTimeout = time.Second
	// FlagTimeout timeout between send flags
	FlagTimeout = time.Second
	// ScoreboardRecalcTimeout timeout between update scoreboard
	ScoreboardRecalcTimeout = time.Second
)

func durationToHMS(d time.Duration) string {

	sec := int(d.Seconds())

	var h, m, s int

	h = sec / 60 / 60
	m = (sec / 60) % 60
	s = sec % 60

	return fmt.Sprintf("%02d:%02d:%02d", h, m, s)
}

func getInfo() string {

	var left time.Duration
	var btnType string

	now := time.Now()

	if now.Before(gameShim.Start) {

		contestStatus = contestNotStarted
		left = gameShim.Start.Sub(now)
		btnType = "stop"

	} else if now.Before(gameShim.End) {

		contestStatus = contestRunning
		left = gameShim.End.Sub(now)
		btnType = "run"

	} else {
		contestStatus = contestCompleted
		left = 0
		btnType = "stop"
	}

	info := fmt.Sprintf(`<span id="game_status-%s">contest %s</span>`,
		btnType, contestStatus)

	if left != 0 {
		info += fmt.Sprintf(`<span id="timer">Left %s</span>`,
			durationToHMS(left))
	}

	return info
}

func infoHandler(ws *websocket.Conn) {

	defer ws.Close()
	for {
		_, err := fmt.Fprint(ws, getInfo())
		if err != nil {
			//log.Println("Socket closed:", err)
			return
		}

		time.Sleep(InfoTimeout)
	}
}

func scoreboardHandler(ws *websocket.Conn) {

	defer ws.Close()

	teamID := getTeamID(ws.Request())

	currentResult := scoreboardHTML(teamID)

	fmt.Fprint(ws, currentResult)

	sendedResult := currentResult

	lastUpdate := time.Now()

	for {
		currentResult := scoreboardHTML(teamID)

		if sendedResult != currentResult ||
			time.Now().After(lastUpdate.Add(time.Minute)) {

			sendedResult = currentResult
			lastUpdate = time.Now()

			_, err := fmt.Fprint(ws, currentResult)
			if err != nil {
				//log.Println("Socket closed:", err)
				return
			}
		}

		time.Sleep(ScoreboardTimeout)
	}
}

func scoreboardHTML(teamID int) (result string) {

	result = "
" +
		"#	" + "Team	" +
		"Score	" +
		"
"

	result += "
"

	scores, err := gameShim.Scoreboard()
	if err != nil {
		log.Println("Get scoreboard fail:", err)
		return
	}

	for n, teamScore := range scores {
		if teamScore.ID == teamID {
			result += `<tr class="self-team">`
		} else {
			result += `<tr>`
		}

		result += fmt.Sprintf(
			`<td class="team_index">%d</td>`+
				`<td class="team_name">%s</td>`+
				`<td class="team_score">%d</td></tr>`,
			n+1, teamScore.Name, teamScore.Score)

	}

	result += "
"

	return
}

func scoreboardUpdater(game *game.Game, updateTimeout time.Duration) {

	for {
		time.Sleep(updateTimeout)

		err := game.RecalcScoreboard()
		if err != nil {
			log.Println("Recalc scoreboard fail:", err)
		}
	}
}

func tasksHTML(teamID int) (result string) {

	cats, err := gameShim.Tasks()
	if err != nil {
		log.Println("Get tasks fail:", err)
	}

	for _, cat := range cats {
		result += categoryToHTML(teamID, cat)
	}

	return
}

func tasksHandler(ws *websocket.Conn) {

	defer ws.Close()

	teamID := getTeamID(ws.Request())

	currentTasks := tasksHTML(teamID)

	fmt.Fprint(ws, currentTasks)

	sendedTasks := currentTasks

	lastUpdate := time.Now()

	for {
		currentTasks := tasksHTML(teamID)

		if sendedTasks != currentTasks ||
			time.Now().After(lastUpdate.Add(time.Minute)) {

			sendedTasks = currentTasks
			lastUpdate = time.Now()

			_, err := fmt.Fprint(ws, currentTasks)
			if err != nil {
				//log.Println("Socket closed:", err)
				return
			}
		}

		time.Sleep(TasksTimeout)
	}
}

func taskHandler(w http.ResponseWriter, r *http.Request) {

	id, err := strconv.Atoi(r.URL.Query().Get("id"))
	if err != nil {
		log.Println("Atoi fail:", err)
		http.Redirect(w, r, "/", 307)
		return
	}

	cats, err := gameShim.Tasks()
	if err != nil {
		log.Println("Get tasks fail:", err)
		http.Redirect(w, r, "/", 307)
		return
	}

	task := game.TaskInfo{ID: id, Opened: false}

	for _, c := range cats {
		for _, t := range c.TasksInfo {
			if t.ID == id {
				task = t
				break
			}
		}
	}

	if !task.Opened {
		// Try to see closed task -> gtfo
		http.Redirect(w, r, "/", 307)
		return
	}

	teamID := getTeamID(r)

	flagSubmitFormat := `<br>` +
		`<form class="input-group" action="/flag?id=%d" method="post">` +
		`<input class="form-control float-left" name="flag" value="" placeholder="Flag">` +
		`<span class="input-group-btn">` +
		`<button class="btn btn-submit">Submit</button>` +
		`</span>` +
		`</form>`

	var submitForm string
	if taskSolvedBy(task, teamID) {
		submitForm = "Already solved"
	} else {
		submitForm = fmt.Sprintf(flagSubmitFormat, task.ID)
	}

	fmt.Fprintf(w, `<!DOCTYPE html>

    <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">

    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <link rel="shortcut icon" href="images/favicon.png" type="image/png">
    <title>Juniors CTF</title>

    <link rel="stylesheet" href="css/style.css" class="--apng-checked">

Scoreboard

Tasks

News

Sponsors

%s

        %s

        %s

          %s

`, task.Name, task.Desc, task.Author, submitForm)
}

func flagHandler(w http.ResponseWriter, r *http.Request) {

	if r.Method != "POST" {
		http.Redirect(w, r, "/", 307)
		return
	}

	taskID, err := strconv.Atoi(r.URL.Query().Get("id"))
	if err != nil {
		log.Println("Atoi fail:", err)
		http.Redirect(w, r, "/", 307)
		return
	}

	flag := r.FormValue("flag")

	teamID := getTeamID(r)

	solved, err := gameShim.Solve(teamID, taskID, flag)
	if err != nil {
		solved = false
	}

	var solvedMsg string
	if solved {
		solvedMsg = `<div class="flag_status solved">Solved</div>`
	} else {
		solvedMsg = `<div class="flag_status invalid">Invalid flag</div>`
	}

	log.Printf("Team ID: %d, Task ID: %d, Flag: %s, Result: %s\n",
		teamID, taskID, flag, solvedMsg)

	time.Sleep(FlagTimeout)

	fmt.Fprintf(w, `<!DOCTYPE html>

    <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">

    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <link rel="shortcut icon" href="images/favicon.png" type="image/png">
    <title>Juniors CTF</title>

    <link rel="stylesheet" href="css/style.css" class="--apng-checked">

Scoreboard

Tasks

News

Sponsors

%s

`, solvedMsg)
}

func handleStaticFile(pattern, file string) {
	http.HandleFunc(pattern,
		func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, file)
		})
}

func handleStaticFileSimple(file, wwwPath string) {
	handleStaticFile(file, wwwPath+file)
}

// Scoreboard implements web scoreboard
func Scoreboard(database *sql.DB, game *game.Game,
	wwwPath, tmpltsPath, addr string) (err error) {

	contestStatus = contestStateNotAvailable
	gameShim = game
	templatePath = tmpltsPath

	go scoreboardUpdater(game, ScoreboardRecalcTimeout)

	// Static files
	handleStaticFileSimple("/css/style.css", wwwPath)
	handleStaticFileSimple("/js/scoreboard.js", wwwPath)
	handleStaticFileSimple("/js/tasks.js", wwwPath)
	handleStaticFileSimple("/news.html", wwwPath)
	handleStaticFileSimple("/sponsors.html", wwwPath)
	handleStaticFileSimple("/images/bg.jpg", wwwPath)
	handleStaticFileSimple("/images/favicon.ico", wwwPath)
	handleStaticFileSimple("/images/favicon.png", wwwPath)
	handleStaticFileSimple("/images/401.jpg", wwwPath)
	handleStaticFileSimple("/images/juniors_ctf_txt.png", wwwPath)
	handleStaticFileSimple("/auth.html", wwwPath)

	// Get
	http.Handle("/", authorized(database, http.HandlerFunc(staticScoreboard)))
	http.Handle("/index.html", authorized(database, http.HandlerFunc(staticScoreboard)))
	http.Handle("/tasks.html", authorized(database, http.HandlerFunc(staticTasks)))
	http.Handle("/logout", authorized(database, http.HandlerFunc(logoutHandler)))

	// Websocket
	http.Handle("/scoreboard", authorized(database, websocket.Handler(scoreboardHandler)))
	http.Handle("/info", authorized(database, websocket.Handler(infoHandler)))
	http.Handle("/tasks", authorized(database, websocket.Handler(tasksHandler)))

	// Post
	http.Handle("/task", authorized(database, http.HandlerFunc(taskHandler)))
	http.Handle("/flag", authorized(database, http.HandlerFunc(flagHandler)))

	http.HandleFunc("/auth.php", http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			authHandler(database, w, r)
		}))

	log.Println("Launching scoreboard at", addr)

	return http.ListenAndServe(addr, nil)
}
<|endoftext|>


<commit_before>package concordances

import (
	"errors"
	"fmt"

	cmneo4j "github.com/Financial-Times/cm-neo4j-driver"
	"github.com/Financial-Times/neo-model-utils-go/mapper"
)

// Driver interface
type Driver interface {
	ReadByConceptID(ids []string) (concordances Concordances, found bool, err error)
	ReadByAuthority(authority string, ids []string) (concordances Concordances, found bool, err error)
	CheckConnectivity() error
}

// CypherDriver struct
type CypherDriver struct {
	driver *cmneo4j.Driver
	env    string
}

//NewCypherDriver instantiate driver
func NewCypherDriver(driver *cmneo4j.Driver, env string) CypherDriver {
	return CypherDriver{driver, env}
}

// CheckConnectivity tests neo4j by running a simple cypher query
func (cd CypherDriver) CheckConnectivity() error {
	return cd.driver.VerifyConnectivity()
}

func (cd CypherDriver) ReadByConceptID(identifiers []string) (concordances Concordances, found bool, err error) {
	var results []neoReadStruct
	query := &cmneo4j.Query{
		Cypher: `
		MATCH (p:Thing)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		MATCH (canonical)<-[:EQUIVALENT_TO]-(leafNode:Thing)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, leafNode.authority as authority, leafNode.authorityValue as authorityValue
		UNION ALL

		MATCH (p:Thing)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		WHERE exists(canonical.leiCode)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, 'LEI' as authority, canonical.leiCode as authorityValue
		UNION ALL

		MATCH (p:Location)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		WHERE exists(canonical.iso31661)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, 'ISO-3166-1' as authority, canonical.iso31661 as authorityValue
		UNION ALL

		MATCH (p:NAICSIndustryClassification)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		WHERE exists(canonical.industryIdentifier)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, 'NAICS' as authority, canonical.industryIdentifier as authorityValue
		UNION ALL

		MATCH (p:Thing)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		MATCH (canonical)<-[:EQUIVALENT_TO]-(leafNode:Thing)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, 'UPP' as authority, leafNode.uuid as authorityValue
        `,
		Params: map[string]interface{}{"identifiers": identifiers},
		Result: &results,
	}

	err = cd.driver.Read(query)
	if errors.Is(err, cmneo4j.ErrNoResultsFound) {
		return Concordances{}, false, nil
	}
	if err != nil {
		return Concordances{}, false, fmt.Errorf("error accessing Concordance datastore for identifier %v: %w", identifiers, err)
	}

	concordances = Concordances{
		Concordance: []Concordance{},
	}

	return processCypherQueryToConcordances(cd, query, results)

}

func (cd CypherDriver) ReadByAuthority(authority string, identifierValues []string) (concordances Concordances, found bool, err error) {
	var results []neoReadStruct

	authorityProperty, found := AuthorityFromURI(authority)
	if !found {
		return Concordances{}, false, nil
	}

	var query *cmneo4j.Query

	if authorityProperty == "UPP" {
		// We need to treat the UPP authority slightly different as it's stored elsewhere.
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (p:Thing)
		WHERE p.uuid IN $authorityValue
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, p.uuid as UUID, 'UPP' as authority, p.uuid as authorityValue`,

			Params: map[string]interface{}{
				"authorityValue": identifierValues,
			},
			Result: &results,
		}
	} else if authorityProperty == "LEI" {
		// We've gotta treat LEI special like as well.
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (p:Concept)
		WHERE p.leiCode IN $authorityValue
		AND exists(p.prefUUID)
		RETURN DISTINCT p.prefUUID AS canonicalUUID, labels(p) AS types, p.uuid as UUID, 'LEI' as authority, p.leiCode as authorityValue`,

			Params: map[string]interface{}{
				"authorityValue": identifierValues,
			},
			Result: &results,
		}
	} else if authorityProperty == "ISO-3166-1" {
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (canonical:Location)
		WHERE canonical.iso31661 IN $authorityValue
		AND exists(canonical.prefUUID)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, canonical.uuid as UUID, 'ISO-3166-1' as authority, canonical.iso31661 as authorityValue
			`,
			Params: map[string]interface{}{
				"authorityValue": identifierValues,
			},
			Result: &results,
		}
	} else if authorityProperty == "NAICS" {
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (canonical:NAICSIndustryClassification)
		WHERE canonical.industryIdentifier IN $authorityValue
		AND exists(canonical.prefUUID)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, canonical.uuid as UUID, 'NAICS' as authority, canonical.industryIdentifier as authorityValue
			`,
			Params: map[string]interface{}{
				"authorityValue": identifierValues,
			},
			Result: &results,
		}
	} else {
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (p:Thing)
		WHERE p.authority = $authority AND p.authorityValue IN $authorityValue
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, p.uuid as UUID, p.authority as authority, p.authorityValue as authorityValue`,

			Params: map[string]interface{}{
				"authorityValue": identifierValues,
				"authority":      authorityProperty,
			},
			Result: &results,
		}
	}

	err = cd.driver.Read(query)
	if errors.Is(err, cmneo4j.ErrNoResultsFound) {
		return Concordances{}, false, nil
	}
	if err != nil {
		return Concordances{}, false, fmt.Errorf("error accessing Concordance datastore for authorityValue %v: %w", identifierValues, err)
	}

	concordances = Concordances{
		Concordance: []Concordance{},
	}

	return processCypherQueryToConcordances(cd, query, results)
}

func processCypherQueryToConcordances(cd CypherDriver, q *cmneo4j.Query, results []neoReadStruct) (concordances Concordances, found bool, err error) {
	err = cd.driver.Read(q)
	if err != nil {
		return Concordances{}, false, fmt.Errorf("error accessing Concordance datastore: %w", err)
	}

	concordances = neoReadStructToConcordances(results, cd.env)

	if errors.Is(err, cmneo4j.ErrNoResultsFound) {
		return Concordances{}, false, nil
	}
	return concordances, true, nil
}

func neoReadStructToConcordances(neo []neoReadStruct, env string) (concordances Concordances) {
	concordances = Concordances{
		Concordance: []Concordance{},
	}
	for _, neoCon := range neo {
		var con = Concordance{}
		var concept = Concept{}

		concept.ID = mapper.IDURL(neoCon.CanonicalUUID)
		concept.APIURL = mapper.APIURL(neoCon.CanonicalUUID, neoCon.Types, env)
		authorityURI, found := AuthorityToURI(neoCon.Authority)
		if !found {
			continue
		}
		con.Identifier = Identifier{Authority: authorityURI, IdentifierValue: neoCon.AuthorityValue}

		con.Concept = concept
		concordances.Concordance = append(concordances.Concordance, con)
	}
	return concordances
}

// Map of authority to URI for the supported concordance IDs
var authorityMap = map[string]string{
	"TME":             "http://api.ft.com/system/FT-TME",
	"FACTSET":         "http://api.ft.com/system/FACTSET",
	"UPP":             "http://api.ft.com/system/UPP",
	"LEI":             "http://api.ft.com/system/LEI",
	"Smartlogic":      "http://api.ft.com/system/SMARTLOGIC",
	"ManagedLocation": "http://api.ft.com/system/MANAGEDLOCATION",
	"ISO-3166-1":      "http://api.ft.com/system/ISO-3166-1",
	"Geonames":        "http://api.ft.com/system/GEONAMES",
	"Wikidata":        "http://api.ft.com/system/WIKIDATA",
	"DBPedia":         "http://api.ft.com/system/DBPEDIA",
	"NAICS":           "http://api.ft.com/system/NAICS",
}

func AuthorityFromURI(uri string) (string, bool) {
	for a, u := range authorityMap {
		if u == uri {
			return a, true
		}
	}
	return "", false
}

func AuthorityToURI(authority string) (string, bool) {
	authorityURI, found := authorityMap[authority]
	return authorityURI, found
}
<commit_msg>Fix error handling<commit_after>package concordances

import (
	"errors"
	"fmt"

	cmneo4j "github.com/Financial-Times/cm-neo4j-driver"
	"github.com/Financial-Times/neo-model-utils-go/mapper"
)

// Driver interface
type Driver interface {
	ReadByConceptID(ids []string) (concordances Concordances, found bool, err error)
	ReadByAuthority(authority string, ids []string) (concordances Concordances, found bool, err error)
	CheckConnectivity() error
}

// CypherDriver struct
type CypherDriver struct {
	driver *cmneo4j.Driver
	env    string
}

//NewCypherDriver instantiate driver
func NewCypherDriver(driver *cmneo4j.Driver, env string) CypherDriver {
	return CypherDriver{driver, env}
}

// CheckConnectivity tests neo4j by running a simple cypher query
func (cd CypherDriver) CheckConnectivity() error {
	return cd.driver.VerifyConnectivity()
}

func (cd CypherDriver) ReadByConceptID(identifiers []string) (concordances Concordances, found bool, err error) {
	var results []neoReadStruct
	query := &cmneo4j.Query{
		Cypher: `
		MATCH (p:Thing)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		MATCH (canonical)<-[:EQUIVALENT_TO]-(leafNode:Thing)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, leafNode.authority as authority, leafNode.authorityValue as authorityValue
		UNION ALL

		MATCH (p:Thing)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		WHERE exists(canonical.leiCode)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, 'LEI' as authority, canonical.leiCode as authorityValue
		UNION ALL

		MATCH (p:Location)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		WHERE exists(canonical.iso31661)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, 'ISO-3166-1' as authority, canonical.iso31661 as authorityValue
		UNION ALL

		MATCH (p:NAICSIndustryClassification)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		WHERE exists(canonical.industryIdentifier)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, 'NAICS' as authority, canonical.industryIdentifier as authorityValue
		UNION ALL

		MATCH (p:Thing)
		WHERE p.uuid in $identifiers
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		MATCH (canonical)<-[:EQUIVALENT_TO]-(leafNode:Thing)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, 'UPP' as authority, leafNode.uuid as authorityValue
        `,
		Params: map[string]interface{}{"identifiers": identifiers},
		Result: &results,
	}

	err = cd.driver.Read(query)

	if errors.Is(err, cmneo4j.ErrNoResultsFound) {
		return Concordances{}, false, nil
	}
	if err != nil {
		return Concordances{}, false, fmt.Errorf("error accessing Concordance datastore for identifier %v: %w", identifiers, err)
	}

	concordances = Concordances{
		Concordance: []Concordance{},
	}

	return processCypherQueryToConcordances(cd, query, results)

}

func (cd CypherDriver) ReadByAuthority(authority string, identifierValues []string) (concordances Concordances, found bool, err error) {
	var results []neoReadStruct

	authorityProperty, found := AuthorityFromURI(authority)
	if !found {
		return Concordances{}, false, nil
	}

	var query *cmneo4j.Query

	if authorityProperty == "UPP" {
		// We need to treat the UPP authority slightly different as it's stored elsewhere.
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (p:Thing)
		WHERE p.uuid IN $authorityValue
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, p.uuid as UUID, 'UPP' as authority, p.uuid as authorityValue`,

			Params: map[string]interface{}{
				"authorityValue": identifierValues,
			},
			Result: &results,
		}
	} else if authorityProperty == "LEI" {
		// We've gotta treat LEI special like as well.
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (p:Concept)
		WHERE p.leiCode IN $authorityValue
		AND exists(p.prefUUID)
		RETURN DISTINCT p.prefUUID AS canonicalUUID, labels(p) AS types, p.uuid as UUID, 'LEI' as authority, p.leiCode as authorityValue`,

			Params: map[string]interface{}{
				"authorityValue": identifierValues,
			},
			Result: &results,
		}
	} else if authorityProperty == "ISO-3166-1" {
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (canonical:Location)
		WHERE canonical.iso31661 IN $authorityValue
		AND exists(canonical.prefUUID)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, canonical.uuid as UUID, 'ISO-3166-1' as authority, canonical.iso31661 as authorityValue
			`,
			Params: map[string]interface{}{
				"authorityValue": identifierValues,
			},
			Result: &results,
		}
	} else if authorityProperty == "NAICS" {
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (canonical:NAICSIndustryClassification)
		WHERE canonical.industryIdentifier IN $authorityValue
		AND exists(canonical.prefUUID)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, canonical.uuid as UUID, 'NAICS' as authority, canonical.industryIdentifier as authorityValue
			`,
			Params: map[string]interface{}{
				"authorityValue": identifierValues,
			},
			Result: &results,
		}
	} else {
		query = &cmneo4j.Query{
			Cypher: `
		MATCH (p:Thing)
		WHERE p.authority = $authority AND p.authorityValue IN $authorityValue
		MATCH (p)-[:EQUIVALENT_TO]->(canonical:Concept)
		RETURN DISTINCT canonical.prefUUID AS canonicalUUID, labels(canonical) AS types, p.uuid as UUID, p.authority as authority, p.authorityValue as authorityValue`,

			Params: map[string]interface{}{
				"authorityValue": identifierValues,
				"authority":      authorityProperty,
			},
			Result: &results,
		}
	}

	err = cd.driver.Read(query)
	if errors.Is(err, cmneo4j.ErrNoResultsFound) {
		return Concordances{}, false, nil
	}
	if err != nil {
		return Concordances{}, false, fmt.Errorf("error accessing Concordance datastore for authorityValue %v: %w", identifierValues, err)
	}

	concordances = Concordances{
		Concordance: []Concordance{},
	}

	return processCypherQueryToConcordances(cd, query, results)
}

func processCypherQueryToConcordances(cd CypherDriver, q *cmneo4j.Query, results []neoReadStruct) (concordances Concordances, found bool, err error) {
	err = cd.driver.Read(q)
	if errors.Is(err, cmneo4j.ErrNoResultsFound) {
		return Concordances{}, false, nil
	}

	if err != nil {
		return Concordances{}, false, fmt.Errorf("error accessing Concordance datastore: %w", err)
	}

	concordances = neoReadStructToConcordances(results, cd.env)

	return concordances, true, nil
}

func neoReadStructToConcordances(neo []neoReadStruct, env string) (concordances Concordances) {
	concordances = Concordances{
		Concordance: []Concordance{},
	}
	for _, neoCon := range neo {
		var con = Concordance{}
		var concept = Concept{}

		concept.ID = mapper.IDURL(neoCon.CanonicalUUID)
		concept.APIURL = mapper.APIURL(neoCon.CanonicalUUID, neoCon.Types, env)
		authorityURI, found := AuthorityToURI(neoCon.Authority)
		if !found {
			continue
		}
		con.Identifier = Identifier{Authority: authorityURI, IdentifierValue: neoCon.AuthorityValue}

		con.Concept = concept
		concordances.Concordance = append(concordances.Concordance, con)
	}
	return concordances
}

// Map of authority to URI for the supported concordance IDs
var authorityMap = map[string]string{
	"TME":             "http://api.ft.com/system/FT-TME",
	"FACTSET":         "http://api.ft.com/system/FACTSET",
	"UPP":             "http://api.ft.com/system/UPP",
	"LEI":             "http://api.ft.com/system/LEI",
	"Smartlogic":      "http://api.ft.com/system/SMARTLOGIC",
	"ManagedLocation": "http://api.ft.com/system/MANAGEDLOCATION",
	"ISO-3166-1":      "http://api.ft.com/system/ISO-3166-1",
	"Geonames":        "http://api.ft.com/system/GEONAMES",
	"Wikidata":        "http://api.ft.com/system/WIKIDATA",
	"DBPedia":         "http://api.ft.com/system/DBPEDIA",
	"NAICS":           "http://api.ft.com/system/NAICS",
}

func AuthorityFromURI(uri string) (string, bool) {
	for a, u := range authorityMap {
		if u == uri {
			return a, true
		}
	}
	return "", false
}

func AuthorityToURI(authority string) (string, bool) {
	authorityURI, found := authorityMap[authority]
	return authorityURI, found
}
<|endoftext|>