<commit_before>package models

import (
	"errors"
	"time"

	"github.com/koding/bongo"
)

type MessageReply struct {
	// unique identifier of the MessageReply
	Id int64 `json:"id,string"`

	// Id of the interacted message
	MessageId int64 `json:"messageId,string"     sql:"NOT NULL"`

	// Id of the reply
	ReplyId int64 `json:"replyId,string"         sql:"NOT NULL"`

	// Creation of the MessageReply
	CreatedAt time.Time `json:"createdAt"         sql:"NOT NULL"`
}

func (m MessageReply) GetId() int64 {
	return m.Id
}

func (m MessageReply) TableName() string {
	return "api.message_reply"
}

func NewMessageReply() *MessageReply {
	return &MessageReply{}
}

func (m *MessageReply) AfterCreate() {
	bongo.B.AfterCreate(m)
}

func (m *MessageReply) AfterUpdate() {
	bongo.B.AfterUpdate(m)
}

func (m MessageReply) AfterDelete() {
	bongo.B.AfterDelete(m)
}

func (m *MessageReply) ById(id int64) error {
	return bongo.B.ById(m, id)
}

func (m *MessageReply) Create() error {
	return bongo.B.Create(m)
}

func (m *MessageReply) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(m, data, q)
}

func (m *MessageReply) One(q *bongo.Query) error {
	return bongo.B.One(m, m, q)
}

func (m *MessageReply) Delete() error {
	selector := map[string]interface{}{
		"message_id": m.MessageId,
		"reply_id":   m.ReplyId,
	}

	if err := m.One(bongo.NewQS(selector)); err != nil {
		return err
	}

	err := bongo.B.Delete(m)
	if err != nil {
		return err
	}

	return nil
}

func (m *MessageReply) DeleteByOrQuery(messageId int64) error {
	var messageReplies []MessageReply
	query := bongo.B.DB.Table(m.TableName())
	query = query.Where("message_id = ? or reply_id = ?", messageId, messageId)

	if err := query.Find(&messageReplies).Error; err != nil {
		return err
	}

	if messageReplies == nil {
		return nil
	}

	if len(messageReplies) == 0 {
		return nil
	}

	for _, messageReply := range messageReplies {
		err := bongo.B.Delete(messageReply)
		if err != nil {
			return err
		}
	}

	return nil
}

func (m *MessageReply) List(query *Query) ([]ChannelMessage, error) {
	return m.fetchMessages(query)
}

func (m *MessageReply) ListAll() ([]ChannelMessage, error) {
	query := NewQuery()
	query.Limit = 0
	query.Skip = 0
	return m.fetchMessages(query)
}

func (m *MessageReply) fetchMessages(query *Query) ([]ChannelMessage, error) {
	var replies []int64

	if m.MessageId == 0 {
		return nil, errors.New("MessageId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": m.MessageId,
		},
		Pluck:      "reply_id",
		Pagination: *bongo.NewPagination(query.Limit, query.Skip),
		Sort:       map[string]string{"created_at": "DESC"},
	}

	bongoQuery := bongo.B.BuildQuery(m, q)
	if !query.From.IsZero() {
		bongoQuery = bongoQuery.Where("created_at < ?", query.From)
	}

	bongoQuery = bongoQuery.Pluck(q.Pluck, &replies)
	if err := bongoQuery.Error; err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessageReplies, err := parent.FetchByIds(replies)
	if err != nil {
		return nil, err
	}

	return channelMessageReplies, nil
}

func (m *MessageReply) UnreadCount(messageId int64, addedAt time.Time) (int, error) {
	if messageId == 0 {
		return 0, errors.New("MessageId is not set")
	}

	if addedAt.IsZero() {
		return 0, errors.New("Last seen at date is not valid - it is zero")
	}

	return bongo.B.Count(
		m,
		"message_id = ? and created_at > ?",
		messageId,
		addedAt.UTC().Format(time.RFC3339),
	)
}

func (m *MessageReply) Count() (int, error) {
	if m.MessageId == 0 {
		return 0, errors.New("MessageId is not set")
	}

	return bongo.B.Count(m,
		"message_id = ?",
		m.MessageId,
	)
}

func (m *MessageReply) FetchRepliedMessage() (*ChannelMessage, error) {
	parent := NewChannelMessage()

	if m.MessageId != 0 {
		if err := parent.ById(m.MessageId); err != nil {
			return nil, err
		}

		return parent, nil
	}

	if m.ReplyId == 0 {
		return nil, errors.New("ReplyId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"reply_id": m.ReplyId,
		},
	}

	if err := m.One(q); err != nil {
		return nil, err
	}

	if err := parent.ById(m.MessageId); err != nil {
		return nil, err
	}

	return parent, nil
}
<commit_msg>Social:  change function name and implement fetcher function for Reply<commit_after>package models

import (
	"errors"
	"time"

	"github.com/koding/bongo"
)

type MessageReply struct {
	// unique identifier of the MessageReply
	Id int64 `json:"id,string"`

	// Id of the interacted message
	MessageId int64 `json:"messageId,string"     sql:"NOT NULL"`

	// Id of the reply
	ReplyId int64 `json:"replyId,string"         sql:"NOT NULL"`

	// Creation of the MessageReply
	CreatedAt time.Time `json:"createdAt"         sql:"NOT NULL"`
}

func (m MessageReply) GetId() int64 {
	return m.Id
}

func (m MessageReply) TableName() string {
	return "api.message_reply"
}

func NewMessageReply() *MessageReply {
	return &MessageReply{}
}

func (m *MessageReply) AfterCreate() {
	bongo.B.AfterCreate(m)
}

func (m *MessageReply) AfterUpdate() {
	bongo.B.AfterUpdate(m)
}

func (m MessageReply) AfterDelete() {
	bongo.B.AfterDelete(m)
}

func (m *MessageReply) ById(id int64) error {
	return bongo.B.ById(m, id)
}

func (m *MessageReply) Create() error {
	return bongo.B.Create(m)
}

func (m *MessageReply) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(m, data, q)
}

func (m *MessageReply) One(q *bongo.Query) error {
	return bongo.B.One(m, m, q)
}

func (m *MessageReply) Delete() error {
	selector := map[string]interface{}{
		"message_id": m.MessageId,
		"reply_id":   m.ReplyId,
	}

	if err := m.One(bongo.NewQS(selector)); err != nil {
		return err
	}

	err := bongo.B.Delete(m)
	if err != nil {
		return err
	}

	return nil
}

func (m *MessageReply) DeleteByOrQuery(messageId int64) error {
	var messageReplies []MessageReply
	query := bongo.B.DB.Table(m.TableName())
	query = query.Where("message_id = ? or reply_id = ?", messageId, messageId)

	if err := query.Find(&messageReplies).Error; err != nil {
		return err
	}

	if messageReplies == nil {
		return nil
	}

	if len(messageReplies) == 0 {
		return nil
	}

	for _, messageReply := range messageReplies {
		err := bongo.B.Delete(messageReply)
		if err != nil {
			return err
		}
	}

	return nil
}

func (m *MessageReply) List(query *Query) ([]ChannelMessage, error) {
	return m.fetchMessages(query)
}

func (m *MessageReply) ListAll() ([]ChannelMessage, error) {
	query := NewQuery()
	query.Limit = 0
	query.Skip = 0
	return m.fetchMessages(query)
}

func (m *MessageReply) fetchMessages(query *Query) ([]ChannelMessage, error) {
	var replies []int64

	if m.MessageId == 0 {
		return nil, errors.New("MessageId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": m.MessageId,
		},
		Pluck:      "reply_id",
		Pagination: *bongo.NewPagination(query.Limit, query.Skip),
		Sort:       map[string]string{"created_at": "DESC"},
	}

	bongoQuery := bongo.B.BuildQuery(m, q)
	if !query.From.IsZero() {
		bongoQuery = bongoQuery.Where("created_at < ?", query.From)
	}

	bongoQuery = bongoQuery.Pluck(q.Pluck, &replies)
	if err := bongoQuery.Error; err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessageReplies, err := parent.FetchByIds(replies)
	if err != nil {
		return nil, err
	}

	return channelMessageReplies, nil
}

func (m *MessageReply) UnreadCount(messageId int64, addedAt time.Time) (int, error) {
	if messageId == 0 {
		return 0, errors.New("MessageId is not set")
	}

	if addedAt.IsZero() {
		return 0, errors.New("Last seen at date is not valid - it is zero")
	}

	return bongo.B.Count(
		m,
		"message_id = ? and created_at > ?",
		messageId,
		addedAt.UTC().Format(time.RFC3339),
	)
}

func (m *MessageReply) Count() (int, error) {
	if m.MessageId == 0 {
		return 0, errors.New("MessageId is not set")
	}

	return bongo.B.Count(m,
		"message_id = ?",
		m.MessageId,
	)
}

func (m *MessageReply) FetchParent() (*ChannelMessage, error) {
	parent := NewChannelMessage()

	if m.ReplyId != 0 {
		if err := parent.ById(m.MessageId); err != nil {
			return nil, err
		}

		return parent, nil
	}

	if m.ReplyId == 0 {
		return nil, errors.New("ReplyId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"reply_id": m.ReplyId,
		},
	}

	if err := m.One(q); err != nil {
		return nil, err
	}

	if err := parent.ById(m.MessageId); err != nil {
		return nil, err
	}

	return parent, nil
}

func (m *MessageReply) FetchReply() (*ChannelMessage, error) {
	reply := NewChannelMessage()

	replyId, err := m.getReplyId()
	if err != nil {
		return nil, err
	}

	if err := reply.ById(replyId); err != nil {
		return nil, err
	}

	return reply, nil
}

func (m *MessageReply) getReplyId() (int64, error) {
	if m.Id == 0 && m.ReplyId == 0 {
		return 0, errors.New("Required ids are not set")
	}

	if m.ReplyId != 0 {
		return m.ReplyId, nil
	}

	if m.Id == 0 {
		// shouldnt come here
		return 0, errors.New("Couldnt fetch replyId")
	}

	if err := m.ById(m.Id); err != nil {
		return 0, err
	}

	return m.ReplyId, nil
}
<|endoftext|>


<commit_before>package tcclient

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-client-go/text"
)

type Credentials struct {
	// Client ID required by Hawk
	ClientId string
	// Access Token required by Hawk
	AccessToken string
	// Certificate for temporary credentials
	Certificate string
	// AuthorizedScopes if set to nil, is ignored. Otherwise, it should be a
	// subset of the scopes that the ClientId already has, and restricts the
	// Credentials to only having these scopes. This is useful when performing
	// actions on behalf of a client which has more restricted scopes. Setting
	// to nil is not the same as setting to an empty array. If AuthorizedScopes
	// is set to an empty array rather than nil, this is equivalent to having
	// no scopes at all.
	// See http://docs.taskcluster.net/auth/authorized-scopes
	AuthorizedScopes []string
}

func (creds *Credentials) String() string {
	return fmt.Sprintf(
		"ClientId: %q\nAccessToken: %q\nCertificate: %q\nAuthorizedScopes: %q",
		creds.ClientId,
		text.StarOut(creds.AccessToken),
		text.StarOut(creds.Certificate),
		creds.AuthorizedScopes,
	)
}

// The entry point into all the functionality in this package is to create a
// ConnectionData object. It contains authentication credentials, and a service
// endpoint, which are required for all HTTP operations.
type ConnectionData struct {
	Credentials *Credentials
	// The URL of the API endpoint to hit.
	// Use "https://auth.taskcluster.net/v1" for production.
	// Please note calling auth.New(clientId string, accessToken string) is an
	// alternative way to create an Auth object with BaseURL set to production.
	BaseURL string
	// Whether authentication is enabled (e.g. set to 'false' when using taskcluster-proxy)
	// Please note calling auth.New(clientId string, accessToken string) is an
	// alternative way to create an Auth object with Authenticate set to true.
	Authenticate bool
}

type Certificate struct {
	Version   int      `json:"version"`
	Scopes    []string `json:"scopes"`
	Start     int64    `json:"start"`
	Expiry    int64    `json:"expiry"`
	Seed      string   `json:"seed"`
	Signature string   `json:"signature"`
	Issuer    string   `json:"issuer,omitempty"`
}

// CreateNamedTemporaryCredentials generates temporary credentials from permanent
// credentials, valid for the given duration, starting immediately.  The
// temporary credentials' scopes must be a subset of the permanent credentials'
// scopes. The duration may not be more than 31 days. Any authorized scopes of
// the permanent credentials will be passed through as authorized scopes to the
// temporary credentials, but will not be restricted via the certificate.
//
// See http://docs.taskcluster.net/auth/temporary-credentials/
func (permaCreds *Credentials) CreateNamedTemporaryCredentials(name string, duration time.Duration, scopes ...string) (tempCreds *Credentials, err error) {
	if duration > 31*24*time.Hour {
		return nil, errors.New("Temporary credentials must expire within 31 days; however a duration of " + duration.String() + " was specified to (*tcclient.ConnectionData).CreateTemporaryCredentials(...) method")
	}

	now := time.Now()
	start := now.Add(time.Minute * -5) // subtract 5 min for clock drift
	expiry := now.Add(duration)

	if permaCreds.ClientId == "" {
		return nil, errors.New("Temporary credentials cannot be created from credentials that have an empty ClientId")
	}
	if permaCreds.AccessToken == "" {
		return nil, errors.New("Temporary credentials cannot be created from credentials that have an empty AccessToken")
	}
	if permaCreds.Certificate != "" {
		return nil, errors.New("Temporary credentials cannot be created from temporary credentials, only from permanent credentials")
	}

	cert := &Certificate{
		Version:   1,
		Scopes:    scopes,
		Start:     start.UnixNano() / 1e6,
		Expiry:    expiry.UnixNano() / 1e6,
		Seed:      slugid.V4() + slugid.V4(),
		Signature: "", // gets set in updateSignature() method below
	}
	// include the issuer iff this is a named credential
	if name != "" {
		cert.Issuer = permaCreds.ClientId
	}

	cert.updateSignature(permaCreds.AccessToken, permaCreds.ClientId, name)

	certBytes, err := json.Marshal(cert)
	if err != nil {
		return
	}

	tempAccessToken, err := generateTemporaryAccessToken(permaCreds.AccessToken, cert.Seed)
	if err != nil {
		return
	}

	tempCreds = &Credentials{
		ClientId:         permaCreds.ClientId,
		AccessToken:      tempAccessToken,
		Certificate:      string(certBytes),
		AuthorizedScopes: permaCreds.AuthorizedScopes,
	}
	if name != "" {
		tempCreds.ClientId = name
	}

	return
}

// CreateTemporaryCredentials is an alias for CreateNamedTemporaryCredentials
// with an empty name.
func (permaCreds *Credentials) CreateTemporaryCredentials(duration time.Duration, scopes ...string) (tempCreds *Credentials, err error) {
	return permaCreds.CreateNamedTemporaryCredentials("", duration, scopes...)
}

func (cert *Certificate) updateSignature(accessToken string, issuer string, name string) (err error) {
	lines := []string{"version:" + strconv.Itoa(cert.Version)}
	// iff this is a named credential, include clientId and issuer
	if name != "" {
		lines = append(lines,
			"clientId:"+name,
			"issuer:"+issuer)
	}
	lines = append(lines,
		"seed:"+cert.Seed,
		"start:"+strconv.FormatInt(cert.Start, 10),
		"expiry:"+strconv.FormatInt(cert.Expiry, 10),
		"scopes:",
	)
	lines = append(lines, cert.Scopes...)
	hash := hmac.New(sha256.New, []byte(accessToken))
	_, err = hash.Write([]byte(strings.Join(lines, "\n")))
	if err != nil {
		return err
	}
	cert.Signature = base64.StdEncoding.EncodeToString(hash.Sum([]byte{}))
	return
}

func generateTemporaryAccessToken(permAccessToken, seed string) (tempAccessToken string, err error) {
	hash := hmac.New(sha256.New, []byte(permAccessToken))
	_, err = hash.Write([]byte(seed))
	if err != nil {
		return "", err
	}
	tempAccessToken = strings.TrimRight(base64.URLEncoding.EncodeToString(hash.Sum([]byte{})), "=")
	return
}

// Attempts to parse the certificate string to return it as an object. If the
// certificate is an empty string (e.g. in the case of permanent credentials)
// then a nil pointer is returned for the certificate. If a certificate has
// been specified but cannot be parsed, an error is returned.
func (creds *Credentials) Cert() (cert *Certificate, err error) {
	if creds.Certificate != "" {
		cert = new(Certificate)
		err = json.Unmarshal([]byte(creds.Certificate), cert)
	}
	return
}
<commit_msg>Updated comments<commit_after>package tcclient

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-client-go/text"
)

type Credentials struct {
	// Client ID required by Hawk
	ClientId string
	// Access Token required by Hawk
	AccessToken string
	// Certificate for temporary credentials
	Certificate string
	// AuthorizedScopes if set to nil, is ignored. Otherwise, it should be a
	// subset of the scopes that the ClientId already has, and restricts the
	// Credentials to only having these scopes. This is useful when performing
	// actions on behalf of a client which has more restricted scopes. Setting
	// to nil is not the same as setting to an empty array. If AuthorizedScopes
	// is set to an empty array rather than nil, this is equivalent to having
	// no scopes at all.
	// See http://docs.taskcluster.net/auth/authorized-scopes
	AuthorizedScopes []string
}

func (creds *Credentials) String() string {
	return fmt.Sprintf(
		"ClientId: %q\nAccessToken: %q\nCertificate: %q\nAuthorizedScopes: %q",
		creds.ClientId,
		text.StarOut(creds.AccessToken),
		text.StarOut(creds.Certificate),
		creds.AuthorizedScopes,
	)
}

// The entry point into all the functionality in this package is to create a
// ConnectionData object. It contains authentication credentials, and a service
// endpoint, which are required for all HTTP operations.
type ConnectionData struct {
	Credentials *Credentials
	// The URL of the API endpoint to hit.
	// Use "https://auth.taskcluster.net/v1" for production.
	// Please note calling auth.New(clientId string, accessToken string) is an
	// alternative way to create an Auth object with BaseURL set to production.
	BaseURL string
	// Whether authentication is enabled (e.g. set to 'false' when using taskcluster-proxy)
	// Please note calling auth.New(clientId string, accessToken string) is an
	// alternative way to create an Auth object with Authenticate set to true.
	Authenticate bool
}

type Certificate struct {
	Version   int      `json:"version"`
	Scopes    []string `json:"scopes"`
	Start     int64    `json:"start"`
	Expiry    int64    `json:"expiry"`
	Seed      string   `json:"seed"`
	Signature string   `json:"signature"`
	Issuer    string   `json:"issuer,omitempty"`
}

// CreateNamedTemporaryCredentials generates temporary credentials from permanent
// credentials, valid for the given duration, starting immediately.  The
// temporary credentials' scopes must be a subset of the permanent credentials'
// scopes. The duration may not be more than 31 days. Any authorized scopes of
// the permanent credentials will be passed through as authorized scopes to the
// temporary credentials, but will not be restricted via the certificate.
//
// See http://docs.taskcluster.net/auth/temporary-credentials/
func (permaCreds *Credentials) CreateNamedTemporaryCredentials(name string, duration time.Duration, scopes ...string) (tempCreds *Credentials, err error) {
	if duration > 31*24*time.Hour {
		return nil, errors.New("Temporary credentials must expire within 31 days; however a duration of " + duration.String() + " was specified to (*tcclient.ConnectionData).CreateTemporaryCredentials(...) method")
	}

	now := time.Now()
	start := now.Add(time.Minute * -5) // subtract 5 min for clock drift
	expiry := now.Add(duration)

	if permaCreds.ClientId == "" {
		return nil, errors.New("Temporary credentials cannot be created from credentials that have an empty ClientId")
	}
	if permaCreds.AccessToken == "" {
		return nil, errors.New("Temporary credentials cannot be created from credentials that have an empty AccessToken")
	}
	if permaCreds.Certificate != "" {
		return nil, errors.New("Temporary credentials cannot be created from temporary credentials, only from permanent credentials")
	}

	cert := &Certificate{
		Version:   1,
		Scopes:    scopes,
		Start:     start.UnixNano() / 1e6,
		Expiry:    expiry.UnixNano() / 1e6,
		Seed:      slugid.V4() + slugid.V4(),
		Signature: "", // gets set in updateSignature() method below
	}
	// include the issuer iff this is a named credential
	if name != "" {
		cert.Issuer = permaCreds.ClientId
	}

	cert.updateSignature(permaCreds.AccessToken, permaCreds.ClientId, name)

	certBytes, err := json.Marshal(cert)
	if err != nil {
		return
	}

	tempAccessToken, err := generateTemporaryAccessToken(permaCreds.AccessToken, cert.Seed)
	if err != nil {
		return
	}

	tempCreds = &Credentials{
		ClientId:         permaCreds.ClientId,
		AccessToken:      tempAccessToken,
		Certificate:      string(certBytes),
		AuthorizedScopes: permaCreds.AuthorizedScopes,
	}
	if name != "" {
		tempCreds.ClientId = name
	}

	return
}

// CreateTemporaryCredentials is an alias for CreateNamedTemporaryCredentials
// with an empty name.
func (permaCreds *Credentials) CreateTemporaryCredentials(duration time.Duration, scopes ...string) (tempCreds *Credentials, err error) {
	return permaCreds.CreateNamedTemporaryCredentials("", duration, scopes...)
}

func (cert *Certificate) updateSignature(accessToken string, issuer string, name string) (err error) {
	lines := []string{"version:" + strconv.Itoa(cert.Version)}
	// iff this is a named credential, include clientId and issuer
	if name != "" {
		lines = append(lines,
			"clientId:"+name,
			"issuer:"+issuer)
	}
	lines = append(lines,
		"seed:"+cert.Seed,
		"start:"+strconv.FormatInt(cert.Start, 10),
		"expiry:"+strconv.FormatInt(cert.Expiry, 10),
		"scopes:",
	)
	lines = append(lines, cert.Scopes...)
	hash := hmac.New(sha256.New, []byte(accessToken))
	_, err = hash.Write([]byte(strings.Join(lines, "\n")))
	if err != nil {
		return err
	}
	cert.Signature = base64.StdEncoding.EncodeToString(hash.Sum([]byte{}))
	return
}

func generateTemporaryAccessToken(permAccessToken, seed string) (tempAccessToken string, err error) {
	hash := hmac.New(sha256.New, []byte(permAccessToken))
	_, err = hash.Write([]byte(seed))
	if err != nil {
		return "", err
	}
	tempAccessToken = strings.TrimRight(base64.URLEncoding.EncodeToString(hash.Sum([]byte{})), "=")
	return
}

// Attempts to parse the certificate string to return it as an object. If the
// certificate is an empty string (e.g. in the case of permanent credentials)
// then a nil pointer is returned for the certificate. If a certificate has
// been specified but cannot be parsed, an error is returned, and cert is an
// empty certificate (rather than nil).
func (creds *Credentials) Cert() (cert *Certificate, err error) {
	if creds.Certificate != "" {
		cert = new(Certificate)
		err = json.Unmarshal([]byte(creds.Certificate), cert)
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2013 Matthew Baird
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//     http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package elastigo

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http/httputil"
	"time"
)

func (c *Conn) DoCommand(method string, url string, args map[string]interface{}, data interface{}) ([]byte, error) {
	var response map[string]interface{}
	var body []byte
	var httpStatusCode int

	query, err := Escape(args)
	if err != nil {
		return nil, err
	}
	req, err := c.NewRequest(method, url, query)
	if err != nil {
		return body, err
	}

	if data != nil {
		switch v := data.(type) {
		case string:
			req.SetBodyString(v)
		case io.Reader:
			req.SetBody(v)
		case []byte:
			req.SetBodyBytes(v)
		default:
			err = req.SetBodyJson(v)
			if err != nil {
				return body, err
			}
		}
	}

	// uncomment this to print out the request that hits the wire
	reqbuf, err := httputil.DumpRequest(req.Request, true)
	log.Println(fmt.Sprintf("\n========= req:\nURL: %s\n%s", req.URL, bytes.NewBuffer(reqbuf).String()))

	// Copy request body for tracer
	if c.RequestTracer != nil {
		requestBody, err := ioutil.ReadAll(req.Body)
		if err != nil {
			return body, err
		}

		req.SetBody(bytes.NewReader(requestBody))
		c.RequestTracer(req.Method, req.URL.String(), string(requestBody))
	}

	httpStatusCode, body, err = req.Do(&response)
	if err != nil {
		return body, err
	}
	if httpStatusCode > 304 {

		jsonErr := json.Unmarshal(body, &response)
		if jsonErr == nil {
			if res_err, ok := response["error"]; ok {
				status, _ := response["status"]
				return body, ESError{time.Now(), fmt.Sprintf("Error [%s] Status [%v]", res_err, status), httpStatusCode}
			}
		}
		return body, jsonErr
	}
	return body, nil
}

// ESError is an error implementation that includes a time, message, and code.
type ESError struct {
	When time.Time
	What string
	Code int
}

func (e ESError) Error() string {
	return fmt.Sprintf("%v: %v [%v]", e.When, e.What, e.Code)
}

// Exists allows the caller to check for the existance of a document using HEAD
// This appears to be broken in the current version of elasticsearch 0.19.10, currently
// returning nothing
func (c *Conn) Exists(index string, _type string, id string, args map[string]interface{}) (BaseResponse, error) {
	var response map[string]interface{}
	var body []byte
	var url string
	var retval BaseResponse
	var httpStatusCode int

	query, err := Escape(args)
	if err != nil {
		return retval, err
	}

	if len(_type) > 0 {
		url = fmt.Sprintf("/%s/%s/%s", index, _type, id)
	} else {
		url = fmt.Sprintf("/%s/%s", index, id)
	}
	req, err := c.NewRequest("HEAD", url, query)
	if err != nil {
		// some sort of generic error handler
	}
	httpStatusCode, body, err = req.Do(&response)
	if httpStatusCode > 304 {
		if error, ok := response["error"]; ok {
			status, _ := response["status"]
			log.Printf("Error: %v (%v)\n", error, status)
		}
	} else {
		// marshall into json
		jsonErr := json.Unmarshal(body, &retval)
		if jsonErr != nil {
			log.Println(jsonErr)
		}
	}
	return retval, err
}
<commit_msg>Comment out debug prints<commit_after>// Copyright 2013 Matthew Baird
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//     http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package elastigo

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http/httputil"
	"time"
)

func (c *Conn) DoCommand(method string, url string, args map[string]interface{}, data interface{}) ([]byte, error) {
	var response map[string]interface{}
	var body []byte
	var httpStatusCode int

	query, err := Escape(args)
	if err != nil {
		return nil, err
	}
	req, err := c.NewRequest(method, url, query)
	if err != nil {
		return body, err
	}

	if data != nil {
		switch v := data.(type) {
		case string:
			req.SetBodyString(v)
		case io.Reader:
			req.SetBody(v)
		case []byte:
			req.SetBodyBytes(v)
		default:
			err = req.SetBodyJson(v)
			if err != nil {
				return body, err
			}
		}
	}

	// uncomment this to print out the request that hits the wire
	//reqbuf, err := httputil.DumpRequest(req.Request, true)
	//log.Println(fmt.Sprintf("\n========= req:\nURL: %s\n%s", req.URL, bytes.NewBuffer(reqbuf).String()))

	// Copy request body for tracer
	if c.RequestTracer != nil {
		requestBody, err := ioutil.ReadAll(req.Body)
		if err != nil {
			return body, err
		}

		req.SetBody(bytes.NewReader(requestBody))
		c.RequestTracer(req.Method, req.URL.String(), string(requestBody))
	}

	httpStatusCode, body, err = req.Do(&response)
	if err != nil {
		return body, err
	}
	if httpStatusCode > 304 {

		jsonErr := json.Unmarshal(body, &response)
		if jsonErr == nil {
			if res_err, ok := response["error"]; ok {
				status, _ := response["status"]
				return body, ESError{time.Now(), fmt.Sprintf("Error [%s] Status [%v]", res_err, status), httpStatusCode}
			}
		}
		return body, jsonErr
	}
	return body, nil
}

// ESError is an error implementation that includes a time, message, and code.
type ESError struct {
	When time.Time
	What string
	Code int
}

func (e ESError) Error() string {
	return fmt.Sprintf("%v: %v [%v]", e.When, e.What, e.Code)
}

// Exists allows the caller to check for the existance of a document using HEAD
// This appears to be broken in the current version of elasticsearch 0.19.10, currently
// returning nothing
func (c *Conn) Exists(index string, _type string, id string, args map[string]interface{}) (BaseResponse, error) {
	var response map[string]interface{}
	var body []byte
	var url string
	var retval BaseResponse
	var httpStatusCode int

	query, err := Escape(args)
	if err != nil {
		return retval, err
	}

	if len(_type) > 0 {
		url = fmt.Sprintf("/%s/%s/%s", index, _type, id)
	} else {
		url = fmt.Sprintf("/%s/%s", index, id)
	}
	req, err := c.NewRequest("HEAD", url, query)
	if err != nil {
		// some sort of generic error handler
	}
	httpStatusCode, body, err = req.Do(&response)
	if httpStatusCode > 304 {
		if error, ok := response["error"]; ok {
			status, _ := response["status"]
			log.Printf("Error: %v (%v)\n", error, status)
		}
	} else {
		// marshall into json
		jsonErr := json.Unmarshal(body, &retval)
		if jsonErr != nil {
			log.Println(jsonErr)
		}
	}
	return retval, err
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"io"
	"log"
	"os"

	"github.com/miekg/dns"
)

var proxyPassAddr = "8.8.8.8:domain"

var blocked = []string{}

func init() {
	file, openErr := os.Open("hosts")
	if nil != openErr {
		panic(openErr)
	}
	defer file.Close()
	buffer := bufio.NewReader(file)
	for {
		line, readErr := buffer.ReadString('\n')
		if readErr == io.EOF {
			break
		} else if nil != readErr {
			panic(readErr)
		}

		host := line[0 : len(line)-1]
		blocked = append(blocked, host)
	}
}

func isBlocked(requestMesage *dns.Msg) bool {
	if 1 != len(requestMesage.Question) {
		log.Printf("Can not process message with multiple questions")
		return false
	}

	question := requestMesage.Question[0]
	for _, name := range blocked {
		if dns.Fqdn(name) == question.Name {
			return true
		}
	}
	return false
}

func proxyRequest(writer dns.ResponseWriter, requestMessage *dns.Msg) {
	responseMessage, exchangeErr := dns.Exchange(requestMessage, proxyPassAddr)

	if nil == exchangeErr {
		log.Printf("Response message: %+v\n", responseMessage)
		writeResponse(writer, responseMessage)
	} else {
		log.Printf("Exchange error: %s\n", exchangeErr)

		errorMessage := new(dns.Msg)
		errorMessage.SetRcode(requestMessage, dns.RcodeServerFailure)
		log.Printf("Error message: %+v", errorMessage)

		writeResponse(writer, errorMessage)
	}
}

func blockedRequest(writer dns.ResponseWriter, requestMessage *dns.Msg) {
	responseMessage := new(dns.Msg)
	responseMessage.SetRcode(requestMessage, dns.RcodeRefused)
	log.Printf("Block response: %+v\n", responseMessage)
	writeResponse(writer, responseMessage)
}

func handler(writer dns.ResponseWriter, requestMessage *dns.Msg) {
	log.Printf("Query message: %+v\n", requestMessage)

	if isBlocked(requestMessage) {
		blockedRequest(writer, requestMessage)
	} else {
		proxyRequest(writer, requestMessage)
	}
}

func writeResponse(writer dns.ResponseWriter, message *dns.Msg) {
	if writeErr := writer.WriteMsg(message); nil == writeErr {
		log.Printf("Writer success\n")
	} else {
		log.Printf("Writer error: %s\n", writeErr)
	}
}

func main() {
	server := &dns.Server{
		Addr:    ":1053",
		Net:     "udp",
		Handler: dns.HandlerFunc(handler),
	}

	if serverErr := server.ListenAndServe(); nil != serverErr {
		log.Fatal(serverErr)
	}
}
<commit_msg>Configure via command line options<commit_after>package main

import (
	"bufio"
	"flag"
	"io"
	"log"
	"os"

	"github.com/miekg/dns"
)

var dnsServer string
var listen string
var hostsPath string
var blocked = []string{}

func init() {
	flag.StringVar(&dnsServer, "dns-server", "192.168.1.1:domain",
		"DNS server for proxy not blocked queries")
	flag.StringVar(&listen, "listen", ":domain", "Listen is pair 'ip:port'")
	flag.StringVar(&hostsPath, "hosts-path", "hosts", "Path to hosts file")
	flag.Parse()

	log.Printf("Start with: listen: %s; dns-server: %s; hosts-path: %s\n",
		listen, dnsServer, hostsPath)

	loadBlocked()
}

func loadBlocked() {
	file, openErr := os.Open(hostsPath)
	if nil != openErr {
		panic(openErr)
	}
	defer file.Close()
	buffer := bufio.NewReader(file)
	for {
		line, readErr := buffer.ReadString('\n')
		if readErr == io.EOF {
			break
		} else if nil != readErr {
			panic(readErr)
		}

		host := line[0 : len(line)-1]
		blocked = append(blocked, host)
	}
}

func isBlocked(requestMesage *dns.Msg) bool {
	if 1 != len(requestMesage.Question) {
		log.Printf("Can not process message with multiple questions")
		return false
	}

	question := requestMesage.Question[0]
	for _, name := range blocked {
		if dns.Fqdn(name) == question.Name {
			return true
		}
	}
	return false
}

func proxyRequest(writer dns.ResponseWriter, requestMessage *dns.Msg) {
	responseMessage, exchangeErr := dns.Exchange(requestMessage, dnsServer)

	if nil == exchangeErr {
		log.Printf("Response message: %+v\n", responseMessage)
		writeResponse(writer, responseMessage)
	} else {
		log.Printf("Exchange error: %s\n", exchangeErr)

		errorMessage := new(dns.Msg)
		errorMessage.SetRcode(requestMessage, dns.RcodeServerFailure)
		log.Printf("Error message: %+v", errorMessage)

		writeResponse(writer, errorMessage)
	}
}

func blockedRequest(writer dns.ResponseWriter, requestMessage *dns.Msg) {
	responseMessage := new(dns.Msg)
	responseMessage.SetRcode(requestMessage, dns.RcodeRefused)
	log.Printf("Block response: %+v\n", responseMessage)
	writeResponse(writer, responseMessage)
}

func handler(writer dns.ResponseWriter, requestMessage *dns.Msg) {
	log.Printf("Query message: %+v\n", requestMessage)

	if isBlocked(requestMessage) {
		blockedRequest(writer, requestMessage)
	} else {
		proxyRequest(writer, requestMessage)
	}
}

func writeResponse(writer dns.ResponseWriter, message *dns.Msg) {
	if writeErr := writer.WriteMsg(message); nil == writeErr {
		log.Printf("Writer success\n")
	} else {
		log.Printf("Writer error: %s\n", writeErr)
	}
}

func main() {
	server := &dns.Server{
		Addr:    listen,
		Net:     "udp",
		Handler: dns.HandlerFunc(handler),
	}

	if serverErr := server.ListenAndServe(); nil != serverErr {
		log.Fatal(serverErr)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

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

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

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

package images

import (
	"context"
	"fmt"
	"github.com/jetstack/cert-manager/hack/release/pkg/util"
	flag "github.com/spf13/pflag"
	"os/exec"

	"github.com/jetstack/cert-manager/hack/release/pkg/bazel"
	"github.com/jetstack/cert-manager/hack/release/pkg/flags"
	logf "github.com/jetstack/cert-manager/hack/release/pkg/log"
)

var (
	Default = &Plugin{}

	supportedGoArch     = []string{"amd64", "arm64", "arm"}
	supportedComponents = []string{"acmesolver", "controller", "webhook", "cainjector"}
	log                 = logf.Log.WithName("images")
)

type Plugin struct {
	// The list of images to build (e.g. acmesolver, controller, webhook)
	Components []string

	// If true, the built images will be exported to the configured docker
	// daemon when the Build() method is called.
	ExportToDocker bool

	// List of architectures to build images for
	GoArch []string

	// TODO: add GOOS support once the build system supports more than linux

	// built is set to true if Build() has completed successfully
	built bool
}

func (g *Plugin) AddFlags(fs *flag.FlagSet) {
	fs.BoolVar(&g.ExportToDocker, "images.export", false, "if true, images will be exported to the currently configured docker daemon")
	fs.StringSliceVar(&g.Components, "images.components", []string{"acmesolver", "controller", "webhook"}, "the list of components to build images for")
	fs.StringSliceVar(&g.GoArch, "images.goarch", []string{"amd64", "arm64", "arm"}, "list of architectures to build images for")
}

func (g *Plugin) Validate() []error {
	var errs []error

	// validate goarch flag
	for _, a := range g.GoArch {
		valid := false
		for _, sa := range supportedGoArch {
			if a == sa {
				valid = true
				break
			}
		}
		if !valid {
			errs = append(errs, fmt.Errorf("invalid goarch value %q", a))
		}
	}

	// validate components flag
	for _, a := range g.Components {
		valid := false
		for _, sa := range supportedComponents {
			if a == sa {
				valid = true
				break
			}
		}
		if !valid {
			errs = append(errs, fmt.Errorf("invalid component name %q", a))
		}
	}

	return errs
}

func (g *Plugin) InitPublish() []error {
	var errs []error

	return errs
}

func (g *Plugin) Build(ctx context.Context) error {
	_, err := g.build(ctx)
	if err != nil {
		return err
	}

	if g.ExportToDocker {
		log.Info("Exporting docker images to local docker daemon")
		if err := g.exportToDocker(ctx); err != nil {
			return err
		}
	} else {
		log.Info("skipping exporting docker images to docker daemon")
	}

	return nil
}

func (g *Plugin) Publish(ctx context.Context) error {
	log.Info("running publish for image plugin")
	// this case should never be reached, but we check it to be safe
	if !g.built {
		if _, err := g.build(ctx); err != nil {
			return err
		}
	}

	log.Info("pushing images")
	targets := g.generateTargets()
	err := g.pushImages(ctx, targets)
	if err != nil {
		return err
	}

	log.Info("published all docker images")

	return nil
}

func (g *Plugin) Complete() error {
	return nil
}

func (g *Plugin) build(ctx context.Context) (imageTargets, error) {
	targets := g.generateTargets()

	bazelTargets := targets.bazelTargets()
	log := log.WithValues("images", bazelTargets)
	log.Info("building bazel image targets")

	// set the os and arch to linux/amd64
	// whilst we might be building cross-arch binaries, cgo only depends on
	// particular OS settings and not arch, so just by setting this to 'linux'
	// we can fix cross builds on platforms other than linux
	// if we support alternate OS values in future, this will need updating
	// with a call to BuildPlatformE per *OS*.
	err := bazel.Default.BuildPlatformE(ctx, log, "linux", "amd64", bazelTargets...)
	if err != nil {
		return nil, fmt.Errorf("error building docker images (%v): %v", targets, err)
	}

	g.built = true
	return targets, nil
}

func (g *Plugin) exportToDocker(ctx context.Context) error {
	targets := g.generateTargets()
	log.WithValues("images", targets.bazelExportTargets()).Info("exporting images to docker daemon")
	for _, target := range targets.bazelExportTargets() {
		log := log.WithValues("target", target)
		log.Info("exporting image to docker daemon")
		// set the os and arch to linux/amd64
		// whilst we might be building cross-arch binaries, cgo only depends on
		// particular OS settings and not arch, so just by setting this to 'linux'
		// we can fix cross builds on platforms other than linux
		// if we support alternate OS values in future, this will need updating
		// with a call to BuildPlatformE per *OS*.
		err := bazel.Default.RunPlatformE(ctx, log, "linux", "amd64", target)
		if err != nil {
			return fmt.Errorf("error exporting image %q to docker daemon: %v", target, err)
		}
	}

	log.WithValues("images", targets.exportedImageNames()).Info("exported all docker images")

	return nil
}

// generateTargets generates a list of Bazel target names that must be
// built for this invocation of the image builder
func (g *Plugin) generateTargets() imageTargets {
	var targets []imageTarget
	for _, c := range g.Components {
		for _, a := range g.GoArch {
			targets = append(targets, imageTarget{c, "linux", a})
		}
	}
	return targets
}

// pushImages will push the images built for this release to the registry
// TODO: add support for calling container_push targets instead of just 'docker push'
func (p *Plugin) pushImages(ctx context.Context, targets imageTargets) error {
	err := p.exportToDocker(ctx)
	if err != nil {
		return err
	}

	var images []string
	for _, t := range targets {
		images = append(images, t.exportedImageNames()...)
	}

	log.WithValues("images", images).Info("pushing docker images")
	for _, img := range images {
		log := log.WithValues("image", img)
		log.Info("pushing docker image")
		err := util.RunE(log, exec.CommandContext(ctx, "docker", "push", img))
		if err != nil {
			return err
		}
	}

	return nil
}

type imageTargets []imageTarget

func (i imageTargets) bazelTargets() []string {
	out := make([]string, len(i))
	for idx, target := range i {
		out[idx] = target.bazelTarget()
	}
	return out
}

func (i imageTargets) bazelExportTargets() []string {
	out := make([]string, len(i))
	for idx, target := range i {
		out[idx] = target.bazelExportTarget()
	}
	return out
}

func (i imageTargets) exportedImageNames() []string {
	out := make([]string, 0)
	for _, target := range i {
		out = append(out, target.exportedImageNames()...)
	}
	return out
}

type imageTarget struct {
	name, os, arch string
}

func (i imageTarget) bazelTarget() string {
	return fmt.Sprintf("//cmd/%s:image.%s-%s", i.name, i.os, i.arch)
}

func (i imageTarget) bazelExportTarget() string {
	return fmt.Sprintf("//cmd/%s:image.%s-%s.export", i.name, i.os, i.arch)
}

func (i imageTarget) exportedImageNames() []string {
	if i.arch == "amd64" {
		return []string{
			fmt.Sprintf("%s/cert-manager-%s:%s", flags.Default.DockerRepo, i.name, flags.Default.AppVersion),
			fmt.Sprintf("%s/cert-manager-%s:%s", flags.Default.DockerRepo, i.name, flags.Default.GitCommitRef),
		}
	}

	return []string{
		fmt.Sprintf("%s/cert-manager-%s-%s:%s", flags.Default.DockerRepo, i.name, i.arch, flags.Default.AppVersion),
		fmt.Sprintf("%s/cert-manager-%s-%s:%s", flags.Default.DockerRepo, i.name, i.arch, flags.Default.GitCommitRef),
	}
}
<commit_msg>Honour DOCKER_CONFIG variable for docker authentication<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

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

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

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

package images

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path"

	flag "github.com/spf13/pflag"

	"github.com/jetstack/cert-manager/hack/release/pkg/bazel"
	"github.com/jetstack/cert-manager/hack/release/pkg/flags"
	logf "github.com/jetstack/cert-manager/hack/release/pkg/log"
	"github.com/jetstack/cert-manager/hack/release/pkg/util"
)

var (
	Default = &Plugin{}

	supportedGoArch     = []string{"amd64", "arm64", "arm"}
	supportedComponents = []string{"acmesolver", "controller", "webhook", "cainjector"}
	log                 = logf.Log.WithName("images")
)

type Plugin struct {
	// The list of images to build (e.g. acmesolver, controller, webhook)
	Components []string

	// If true, the built images will be exported to the configured docker
	// daemon when the Build() method is called.
	ExportToDocker bool

	// List of architectures to build images for
	GoArch []string

	// DockerConfig is a path to a directory containing a config.json file that
	// is used for Docker authentication
	DockerConfig string

	// TODO: add GOOS support once the build system supports more than linux

	// built is set to true if Build() has completed successfully
	built bool
	// configFileName is computed based on the DockerConfig field
	configFileName string
}

func (g *Plugin) AddFlags(fs *flag.FlagSet) {
	fs.BoolVar(&g.ExportToDocker, "images.export", false, "if true, images will be exported to the currently configured docker daemon")
	fs.StringSliceVar(&g.Components, "images.components", []string{"acmesolver", "controller", "webhook"}, "the list of components to build images for")
	fs.StringSliceVar(&g.GoArch, "images.goarch", []string{"amd64", "arm64", "arm"}, "list of architectures to build images for")
	fs.StringVar(&g.DockerConfig, "images.docker-config", "", "path to a directory containing a docker config.json file used when pushing images")
}

func (g *Plugin) Validate() []error {
	var errs []error

	// validate goarch flag
	for _, a := range g.GoArch {
		valid := false
		for _, sa := range supportedGoArch {
			if a == sa {
				valid = true
				break
			}
		}
		if !valid {
			errs = append(errs, fmt.Errorf("invalid goarch value %q", a))
		}
	}

	// validate components flag
	for _, a := range g.Components {
		valid := false
		for _, sa := range supportedComponents {
			if a == sa {
				valid = true
				break
			}
		}
		if !valid {
			errs = append(errs, fmt.Errorf("invalid component name %q", a))
		}
	}

	return errs
}

func (g *Plugin) InitPublish() []error {
	var errs []error

	if g.DockerConfig != "" {
		// Hardcode the filename to be config.json for backwards compatibility
		// with the old release.sh script.
		configFileName := path.Join(g.DockerConfig, "config.json")
		f, err := os.Stat(configFileName)
		if err != nil {
			return []error{fmt.Errorf("error checking config file: %v", err)}
		}
		if f.IsDir() {
			return []error{fmt.Errorf("docker config.json is not a file")}
		}
		g.configFileName = configFileName
	}

	return errs
}

func (g *Plugin) Build(ctx context.Context) error {
	_, err := g.build(ctx)
	if err != nil {
		return err
	}

	if g.ExportToDocker {
		log.Info("Exporting docker images to local docker daemon")
		if err := g.exportToDocker(ctx); err != nil {
			return err
		}
	} else {
		log.Info("skipping exporting docker images to docker daemon")
	}

	return nil
}

func (g *Plugin) Publish(ctx context.Context) error {
	log.Info("running publish for image plugin")
	// this case should never be reached, but we check it to be safe
	if !g.built {
		if _, err := g.build(ctx); err != nil {
			return err
		}
	}

	log.Info("pushing images")
	targets := g.generateTargets()
	err := g.pushImages(ctx, targets)
	if err != nil {
		return err
	}

	log.Info("published all docker images")

	return nil
}

func (g *Plugin) Complete() error {
	log = log.WithName("default-flags")

	if g.DockerConfig == "" {
		g.DockerConfig = os.Getenv("DOCKER_CONFIG")
		if g.DockerConfig != "" {
			log.Info("set default value", "value", g.DockerConfig)
		}
	}
	return nil
}

func (g *Plugin) build(ctx context.Context) (imageTargets, error) {
	targets := g.generateTargets()

	bazelTargets := targets.bazelTargets()
	log := log.WithValues("images", bazelTargets)
	log.Info("building bazel image targets")

	// set the os and arch to linux/amd64
	// whilst we might be building cross-arch binaries, cgo only depends on
	// particular OS settings and not arch, so just by setting this to 'linux'
	// we can fix cross builds on platforms other than linux
	// if we support alternate OS values in future, this will need updating
	// with a call to BuildPlatformE per *OS*.
	err := bazel.Default.BuildPlatformE(ctx, log, "linux", "amd64", bazelTargets...)
	if err != nil {
		return nil, fmt.Errorf("error building docker images (%v): %v", targets, err)
	}

	g.built = true
	return targets, nil
}

func (g *Plugin) exportToDocker(ctx context.Context) error {
	targets := g.generateTargets()
	log.WithValues("images", targets.bazelExportTargets()).Info("exporting images to docker daemon")
	for _, target := range targets.bazelExportTargets() {
		log := log.WithValues("target", target)
		log.Info("exporting image to docker daemon")
		// set the os and arch to linux/amd64
		// whilst we might be building cross-arch binaries, cgo only depends on
		// particular OS settings and not arch, so just by setting this to 'linux'
		// we can fix cross builds on platforms other than linux
		// if we support alternate OS values in future, this will need updating
		// with a call to BuildPlatformE per *OS*.
		err := bazel.Default.RunPlatformE(ctx, log, "linux", "amd64", target)
		if err != nil {
			return fmt.Errorf("error exporting image %q to docker daemon: %v", target, err)
		}
	}

	log.WithValues("images", targets.exportedImageNames()).Info("exported all docker images")

	return nil
}

// generateTargets generates a list of Bazel target names that must be
// built for this invocation of the image builder
func (g *Plugin) generateTargets() imageTargets {
	var targets []imageTarget
	for _, c := range g.Components {
		for _, a := range g.GoArch {
			targets = append(targets, imageTarget{c, "linux", a})
		}
	}
	return targets
}

// pushImages will push the images built for this release to the registry
// TODO: add support for calling container_push targets instead of just 'docker push'
func (p *Plugin) pushImages(ctx context.Context, targets imageTargets) error {
	err := p.exportToDocker(ctx)
	if err != nil {
		return err
	}

	var images []string
	for _, t := range targets {
		images = append(images, t.exportedImageNames()...)
	}

	log.WithValues("images", images).Info("pushing docker images")
	for _, img := range images {
		log := log.WithValues("image", img)
		log.Info("pushing docker image")
		args := []string{"push", img}
		if p.configFileName != "" {
			args = append(args, "--config", p.configFileName)
		}
		cmd := exec.CommandContext(ctx, "docker", args...)
		err := util.RunE(log, cmd)
		if err != nil {
			return err
		}
	}

	return nil
}

type imageTargets []imageTarget

func (i imageTargets) bazelTargets() []string {
	out := make([]string, len(i))
	for idx, target := range i {
		out[idx] = target.bazelTarget()
	}
	return out
}

func (i imageTargets) bazelExportTargets() []string {
	out := make([]string, len(i))
	for idx, target := range i {
		out[idx] = target.bazelExportTarget()
	}
	return out
}

func (i imageTargets) exportedImageNames() []string {
	out := make([]string, 0)
	for _, target := range i {
		out = append(out, target.exportedImageNames()...)
	}
	return out
}

type imageTarget struct {
	name, os, arch string
}

func (i imageTarget) bazelTarget() string {
	return fmt.Sprintf("//cmd/%s:image.%s-%s", i.name, i.os, i.arch)
}

func (i imageTarget) bazelExportTarget() string {
	return fmt.Sprintf("//cmd/%s:image.%s-%s.export", i.name, i.os, i.arch)
}

func (i imageTarget) exportedImageNames() []string {
	if i.arch == "amd64" {
		return []string{
			fmt.Sprintf("%s/cert-manager-%s:%s", flags.Default.DockerRepo, i.name, flags.Default.AppVersion),
			fmt.Sprintf("%s/cert-manager-%s:%s", flags.Default.DockerRepo, i.name, flags.Default.GitCommitRef),
		}
	}

	return []string{
		fmt.Sprintf("%s/cert-manager-%s-%s:%s", flags.Default.DockerRepo, i.name, i.arch, flags.Default.AppVersion),
		fmt.Sprintf("%s/cert-manager-%s-%s:%s", flags.Default.DockerRepo, i.name, i.arch, flags.Default.GitCommitRef),
	}
}
<|endoftext|>


<commit_before>package shadowsocks

import (
	"bytes"
	"crypto/rand"
	"io"

	"v2ray.com/core/common"
	"v2ray.com/core/common/bitmask"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
)

const (
	Version                               = 1
	RequestOptionOneTimeAuth bitmask.Byte = 0x01
)

var addrParser = protocol.NewAddressParser(
	protocol.AddressFamilyByte(0x01, net.AddressFamilyIPv4),
	protocol.AddressFamilyByte(0x04, net.AddressFamilyIPv6),
	protocol.AddressFamilyByte(0x03, net.AddressFamilyDomain),
	protocol.WithAddressTypeParser(func(b byte) byte {
		return b & 0x0F
	}),
)

// ReadTCPSession reads a Shadowsocks TCP session from the given reader, returns its header and remaining parts.
func ReadTCPSession(user *protocol.MemoryUser, reader io.Reader) (*protocol.RequestHeader, buf.Reader, error) {
	account := user.Account.(*MemoryAccount)

	buffer := buf.New()
	defer buffer.Release()

	ivLen := account.Cipher.IVSize()
	var iv []byte
	if ivLen > 0 {
		if _, err := buffer.ReadFullFrom(reader, ivLen); err != nil {
			return nil, nil, newError("failed to read IV").Base(err)
		}

		iv = append([]byte(nil), buffer.BytesTo(ivLen)...)
	}

	r, err := account.Cipher.NewDecryptionReader(account.Key, iv, reader)
	if err != nil {
		return nil, nil, newError("failed to initialize decoding stream").Base(err).AtError()
	}
	br := &buf.BufferedReader{Reader: r}
	reader = nil

	authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
	request := &protocol.RequestHeader{
		Version: Version,
		User:    user,
		Command: protocol.RequestCommandTCP,
	}

	buffer.Clear()

	addr, port, err := addrParser.ReadAddressPort(buffer, br)
	if err != nil {
		return nil, nil, newError("failed to read address").Base(err)
	}

	request.Address = addr
	request.Port = port

	if !account.Cipher.IsAEAD() {
		if (buffer.Byte(0) & 0x10) == 0x10 {
			request.Option.Set(RequestOptionOneTimeAuth)
		}

		if request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Disabled {
			return nil, nil, newError("rejecting connection with OTA enabled, while server disables OTA")
		}

		if !request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Enabled {
			return nil, nil, newError("rejecting connection with OTA disabled, while server enables OTA")
		}
	}

	if request.Option.Has(RequestOptionOneTimeAuth) {
		actualAuth := make([]byte, AuthSize)
		authenticator.Authenticate(buffer.Bytes())(actualAuth)

		_, err := buffer.ReadFullFrom(br, AuthSize)
		if err != nil {
			return nil, nil, newError("Failed to read OTA").Base(err)
		}

		if !bytes.Equal(actualAuth, buffer.BytesFrom(-AuthSize)) {
			return nil, nil, newError("invalid OTA")
		}
	}

	if request.Address == nil {
		return nil, nil, newError("invalid remote address.")
	}

	var chunkReader buf.Reader
	if request.Option.Has(RequestOptionOneTimeAuth) {
		chunkReader = NewChunkReader(br, NewAuthenticator(ChunkKeyGenerator(iv)))
	} else {
		chunkReader = buf.NewReader(br)
	}

	return request, chunkReader, nil
}

// WriteTCPRequest writes Shadowsocks request into the given writer, and returns a writer for body.
func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (buf.Writer, error) {
	user := request.User
	account := user.Account.(*MemoryAccount)

	if account.Cipher.IsAEAD() {
		request.Option.Clear(RequestOptionOneTimeAuth)
	}

	var iv []byte
	if account.Cipher.IVSize() > 0 {
		iv = make([]byte, account.Cipher.IVSize())
		common.Must2(rand.Read(iv))
		if err := buf.WriteAllBytes(writer, iv); err != nil {
			return nil, newError("failed to write IV")
		}
	}

	w, err := account.Cipher.NewEncryptionWriter(account.Key, iv, writer)
	if err != nil {
		return nil, newError("failed to create encoding stream").Base(err).AtError()
	}

	header := buf.New()

	if err := addrParser.WriteAddressPort(header, request.Address, request.Port); err != nil {
		return nil, newError("failed to write address").Base(err)
	}

	if request.Option.Has(RequestOptionOneTimeAuth) {
		header.SetByte(0, header.Byte(0)|0x10)

		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
		authBuffer := header.Extend(AuthSize)
		authenticator.Authenticate(header.Bytes(), authBuffer)
	}

	if err := w.WriteMultiBuffer(buf.NewMultiBufferValue(header)); err != nil {
		return nil, newError("failed to write header").Base(err)
	}

	var chunkWriter buf.Writer
	if request.Option.Has(RequestOptionOneTimeAuth) {
		chunkWriter = NewChunkWriter(w.(io.Writer), NewAuthenticator(ChunkKeyGenerator(iv)))
	} else {
		chunkWriter = w
	}

	return chunkWriter, nil
}

func ReadTCPResponse(user *protocol.MemoryUser, reader io.Reader) (buf.Reader, error) {
	account := user.Account.(*MemoryAccount)

	var iv []byte
	if account.Cipher.IVSize() > 0 {
		iv = make([]byte, account.Cipher.IVSize())
		if _, err := io.ReadFull(reader, iv); err != nil {
			return nil, newError("failed to read IV").Base(err)
		}
	}

	return account.Cipher.NewDecryptionReader(account.Key, iv, reader)
}

func WriteTCPResponse(request *protocol.RequestHeader, writer io.Writer) (buf.Writer, error) {
	user := request.User
	account := user.Account.(*MemoryAccount)

	var iv []byte
	if account.Cipher.IVSize() > 0 {
		iv = make([]byte, account.Cipher.IVSize())
		common.Must2(rand.Read(iv))
		if err := buf.WriteAllBytes(writer, iv); err != nil {
			return nil, newError("failed to write IV.").Base(err)
		}
	}

	return account.Cipher.NewEncryptionWriter(account.Key, iv, writer)
}

func EncodeUDPPacket(request *protocol.RequestHeader, payload []byte) (*buf.Buffer, error) {
	user := request.User
	account := user.Account.(*MemoryAccount)

	buffer := buf.New()
	ivLen := account.Cipher.IVSize()
	if ivLen > 0 {
		common.Must2(buffer.ReadFullFrom(rand.Reader, ivLen))
	}
	iv := buffer.Bytes()

	if err := addrParser.WriteAddressPort(buffer, request.Address, request.Port); err != nil {
		return nil, newError("failed to write address").Base(err)
	}

	buffer.Write(payload)

	if !account.Cipher.IsAEAD() && request.Option.Has(RequestOptionOneTimeAuth) {
		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
		buffer.SetByte(ivLen, buffer.Byte(ivLen)|0x10)

		authPayload := buffer.BytesFrom(ivLen)
		authBuffer := buffer.Extend(AuthSize)
		authenticator.Authenticate(authPayload, authBuffer)
	}
	if err := account.Cipher.EncodePacket(account.Key, buffer); err != nil {
		return nil, newError("failed to encrypt UDP payload").Base(err)
	}

	return buffer, nil
}

func DecodeUDPPacket(user *protocol.MemoryUser, payload *buf.Buffer) (*protocol.RequestHeader, *buf.Buffer, error) {
	account := user.Account.(*MemoryAccount)

	var iv []byte
	if !account.Cipher.IsAEAD() && account.Cipher.IVSize() > 0 {
		// Keep track of IV as it gets removed from payload in DecodePacket.
		iv = make([]byte, account.Cipher.IVSize())
		copy(iv, payload.BytesTo(account.Cipher.IVSize()))
	}

	if err := account.Cipher.DecodePacket(account.Key, payload); err != nil {
		return nil, nil, newError("failed to decrypt UDP payload").Base(err)
	}

	request := &protocol.RequestHeader{
		Version: Version,
		User:    user,
		Command: protocol.RequestCommandUDP,
	}

	if !account.Cipher.IsAEAD() {
		if (payload.Byte(0) & 0x10) == 0x10 {
			request.Option |= RequestOptionOneTimeAuth
		}

		if request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Disabled {
			return nil, nil, newError("rejecting packet with OTA enabled, while server disables OTA").AtWarning()
		}

		if !request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Enabled {
			return nil, nil, newError("rejecting packet with OTA disabled, while server enables OTA").AtWarning()
		}

		if request.Option.Has(RequestOptionOneTimeAuth) {
			payloadLen := payload.Len() - AuthSize
			authBytes := payload.BytesFrom(payloadLen)

			authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
			actualAuth := make([]byte, AuthSize)
			common.Must2(authenticator.Authenticate(payload.BytesTo(payloadLen))(actualAuth))
			if !bytes.Equal(actualAuth, authBytes) {
				return nil, nil, newError("invalid OTA")
			}

			payload.Resize(0, payloadLen)
		}
	}

	payload.SetByte(0, payload.Byte(0)&0x0F)

	addr, port, err := addrParser.ReadAddressPort(nil, payload)
	if err != nil {
		return nil, nil, newError("failed to parse address").Base(err)
	}

	request.Address = addr
	request.Port = port

	return request, payload, nil
}

type UDPReader struct {
	Reader io.Reader
	User   *protocol.MemoryUser
}

func (v *UDPReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
	buffer := buf.New()
	_, err := buffer.ReadFrom(v.Reader)
	if err != nil {
		buffer.Release()
		return nil, err
	}
	_, payload, err := DecodeUDPPacket(v.User, buffer)
	if err != nil {
		buffer.Release()
		return nil, err
	}
	return buf.NewMultiBufferValue(payload), nil
}

type UDPWriter struct {
	Writer  io.Writer
	Request *protocol.RequestHeader
}

// Write implements io.Writer.
func (w *UDPWriter) Write(payload []byte) (int, error) {
	packet, err := EncodeUDPPacket(w.Request, payload)
	if err != nil {
		return 0, err
	}
	_, err = w.Writer.Write(packet.Bytes())
	packet.Release()
	return len(payload), err
}
<commit_msg>fix build break in shadowsocks<commit_after>package shadowsocks

import (
	"bytes"
	"crypto/rand"
	"io"

	"v2ray.com/core/common"
	"v2ray.com/core/common/bitmask"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
)

const (
	Version                               = 1
	RequestOptionOneTimeAuth bitmask.Byte = 0x01
)

var addrParser = protocol.NewAddressParser(
	protocol.AddressFamilyByte(0x01, net.AddressFamilyIPv4),
	protocol.AddressFamilyByte(0x04, net.AddressFamilyIPv6),
	protocol.AddressFamilyByte(0x03, net.AddressFamilyDomain),
	protocol.WithAddressTypeParser(func(b byte) byte {
		return b & 0x0F
	}),
)

// ReadTCPSession reads a Shadowsocks TCP session from the given reader, returns its header and remaining parts.
func ReadTCPSession(user *protocol.MemoryUser, reader io.Reader) (*protocol.RequestHeader, buf.Reader, error) {
	account := user.Account.(*MemoryAccount)

	buffer := buf.New()
	defer buffer.Release()

	ivLen := account.Cipher.IVSize()
	var iv []byte
	if ivLen > 0 {
		if _, err := buffer.ReadFullFrom(reader, ivLen); err != nil {
			return nil, nil, newError("failed to read IV").Base(err)
		}

		iv = append([]byte(nil), buffer.BytesTo(ivLen)...)
	}

	r, err := account.Cipher.NewDecryptionReader(account.Key, iv, reader)
	if err != nil {
		return nil, nil, newError("failed to initialize decoding stream").Base(err).AtError()
	}
	br := &buf.BufferedReader{Reader: r}
	reader = nil

	authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
	request := &protocol.RequestHeader{
		Version: Version,
		User:    user,
		Command: protocol.RequestCommandTCP,
	}

	buffer.Clear()

	addr, port, err := addrParser.ReadAddressPort(buffer, br)
	if err != nil {
		return nil, nil, newError("failed to read address").Base(err)
	}

	request.Address = addr
	request.Port = port

	if !account.Cipher.IsAEAD() {
		if (buffer.Byte(0) & 0x10) == 0x10 {
			request.Option.Set(RequestOptionOneTimeAuth)
		}

		if request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Disabled {
			return nil, nil, newError("rejecting connection with OTA enabled, while server disables OTA")
		}

		if !request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Enabled {
			return nil, nil, newError("rejecting connection with OTA disabled, while server enables OTA")
		}
	}

	if request.Option.Has(RequestOptionOneTimeAuth) {
		actualAuth := make([]byte, AuthSize)
		authenticator.Authenticate(buffer.Bytes(), actualAuth)

		_, err := buffer.ReadFullFrom(br, AuthSize)
		if err != nil {
			return nil, nil, newError("Failed to read OTA").Base(err)
		}

		if !bytes.Equal(actualAuth, buffer.BytesFrom(-AuthSize)) {
			return nil, nil, newError("invalid OTA")
		}
	}

	if request.Address == nil {
		return nil, nil, newError("invalid remote address.")
	}

	var chunkReader buf.Reader
	if request.Option.Has(RequestOptionOneTimeAuth) {
		chunkReader = NewChunkReader(br, NewAuthenticator(ChunkKeyGenerator(iv)))
	} else {
		chunkReader = buf.NewReader(br)
	}

	return request, chunkReader, nil
}

// WriteTCPRequest writes Shadowsocks request into the given writer, and returns a writer for body.
func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (buf.Writer, error) {
	user := request.User
	account := user.Account.(*MemoryAccount)

	if account.Cipher.IsAEAD() {
		request.Option.Clear(RequestOptionOneTimeAuth)
	}

	var iv []byte
	if account.Cipher.IVSize() > 0 {
		iv = make([]byte, account.Cipher.IVSize())
		common.Must2(rand.Read(iv))
		if err := buf.WriteAllBytes(writer, iv); err != nil {
			return nil, newError("failed to write IV")
		}
	}

	w, err := account.Cipher.NewEncryptionWriter(account.Key, iv, writer)
	if err != nil {
		return nil, newError("failed to create encoding stream").Base(err).AtError()
	}

	header := buf.New()

	if err := addrParser.WriteAddressPort(header, request.Address, request.Port); err != nil {
		return nil, newError("failed to write address").Base(err)
	}

	if request.Option.Has(RequestOptionOneTimeAuth) {
		header.SetByte(0, header.Byte(0)|0x10)

		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
		authBuffer := header.Extend(AuthSize)
		authenticator.Authenticate(header.Bytes(), authBuffer)
	}

	if err := w.WriteMultiBuffer(buf.NewMultiBufferValue(header)); err != nil {
		return nil, newError("failed to write header").Base(err)
	}

	var chunkWriter buf.Writer
	if request.Option.Has(RequestOptionOneTimeAuth) {
		chunkWriter = NewChunkWriter(w.(io.Writer), NewAuthenticator(ChunkKeyGenerator(iv)))
	} else {
		chunkWriter = w
	}

	return chunkWriter, nil
}

func ReadTCPResponse(user *protocol.MemoryUser, reader io.Reader) (buf.Reader, error) {
	account := user.Account.(*MemoryAccount)

	var iv []byte
	if account.Cipher.IVSize() > 0 {
		iv = make([]byte, account.Cipher.IVSize())
		if _, err := io.ReadFull(reader, iv); err != nil {
			return nil, newError("failed to read IV").Base(err)
		}
	}

	return account.Cipher.NewDecryptionReader(account.Key, iv, reader)
}

func WriteTCPResponse(request *protocol.RequestHeader, writer io.Writer) (buf.Writer, error) {
	user := request.User
	account := user.Account.(*MemoryAccount)

	var iv []byte
	if account.Cipher.IVSize() > 0 {
		iv = make([]byte, account.Cipher.IVSize())
		common.Must2(rand.Read(iv))
		if err := buf.WriteAllBytes(writer, iv); err != nil {
			return nil, newError("failed to write IV.").Base(err)
		}
	}

	return account.Cipher.NewEncryptionWriter(account.Key, iv, writer)
}

func EncodeUDPPacket(request *protocol.RequestHeader, payload []byte) (*buf.Buffer, error) {
	user := request.User
	account := user.Account.(*MemoryAccount)

	buffer := buf.New()
	ivLen := account.Cipher.IVSize()
	if ivLen > 0 {
		common.Must2(buffer.ReadFullFrom(rand.Reader, ivLen))
	}
	iv := buffer.Bytes()

	if err := addrParser.WriteAddressPort(buffer, request.Address, request.Port); err != nil {
		return nil, newError("failed to write address").Base(err)
	}

	buffer.Write(payload)

	if !account.Cipher.IsAEAD() && request.Option.Has(RequestOptionOneTimeAuth) {
		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
		buffer.SetByte(ivLen, buffer.Byte(ivLen)|0x10)

		authPayload := buffer.BytesFrom(ivLen)
		authBuffer := buffer.Extend(AuthSize)
		authenticator.Authenticate(authPayload, authBuffer)
	}
	if err := account.Cipher.EncodePacket(account.Key, buffer); err != nil {
		return nil, newError("failed to encrypt UDP payload").Base(err)
	}

	return buffer, nil
}

func DecodeUDPPacket(user *protocol.MemoryUser, payload *buf.Buffer) (*protocol.RequestHeader, *buf.Buffer, error) {
	account := user.Account.(*MemoryAccount)

	var iv []byte
	if !account.Cipher.IsAEAD() && account.Cipher.IVSize() > 0 {
		// Keep track of IV as it gets removed from payload in DecodePacket.
		iv = make([]byte, account.Cipher.IVSize())
		copy(iv, payload.BytesTo(account.Cipher.IVSize()))
	}

	if err := account.Cipher.DecodePacket(account.Key, payload); err != nil {
		return nil, nil, newError("failed to decrypt UDP payload").Base(err)
	}

	request := &protocol.RequestHeader{
		Version: Version,
		User:    user,
		Command: protocol.RequestCommandUDP,
	}

	if !account.Cipher.IsAEAD() {
		if (payload.Byte(0) & 0x10) == 0x10 {
			request.Option |= RequestOptionOneTimeAuth
		}

		if request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Disabled {
			return nil, nil, newError("rejecting packet with OTA enabled, while server disables OTA").AtWarning()
		}

		if !request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Enabled {
			return nil, nil, newError("rejecting packet with OTA disabled, while server enables OTA").AtWarning()
		}

		if request.Option.Has(RequestOptionOneTimeAuth) {
			payloadLen := payload.Len() - AuthSize
			authBytes := payload.BytesFrom(payloadLen)

			authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
			actualAuth := make([]byte, AuthSize)
			authenticator.Authenticate(payload.BytesTo(payloadLen), actualAuth)
			if !bytes.Equal(actualAuth, authBytes) {
				return nil, nil, newError("invalid OTA")
			}

			payload.Resize(0, payloadLen)
		}
	}

	payload.SetByte(0, payload.Byte(0)&0x0F)

	addr, port, err := addrParser.ReadAddressPort(nil, payload)
	if err != nil {
		return nil, nil, newError("failed to parse address").Base(err)
	}

	request.Address = addr
	request.Port = port

	return request, payload, nil
}

type UDPReader struct {
	Reader io.Reader
	User   *protocol.MemoryUser
}

func (v *UDPReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
	buffer := buf.New()
	_, err := buffer.ReadFrom(v.Reader)
	if err != nil {
		buffer.Release()
		return nil, err
	}
	_, payload, err := DecodeUDPPacket(v.User, buffer)
	if err != nil {
		buffer.Release()
		return nil, err
	}
	return buf.NewMultiBufferValue(payload), nil
}

type UDPWriter struct {
	Writer  io.Writer
	Request *protocol.RequestHeader
}

// Write implements io.Writer.
func (w *UDPWriter) Write(payload []byte) (int, error) {
	packet, err := EncodeUDPPacket(w.Request, payload)
	if err != nil {
		return 0, err
	}
	_, err = w.Writer.Write(packet.Bytes())
	packet.Release()
	return len(payload), err
}
<|endoftext|>


<commit_before>package needle

import (
	. "github.com/chrislusf/seaweedfs/weed/storage/types"
	"sort"
	"sync"
)

const (
	batch = 100000
)

type SectionalNeedleId uint32

const SectionalNeedleIdLimit = 1<<32 - 1

type SectionalNeedleValue struct {
	Key    SectionalNeedleId
	Offset Offset `comment:"Volume offset"` //since aligned to 8 bytes, range is 4G*8=32G
	Size   uint32 `comment:"Size of the data portion"`
}

type CompactSection struct {
	sync.RWMutex
	values   []SectionalNeedleValue
	overflow Overflow
	start    NeedleId
	end      NeedleId
	counter  int
}

type Overflow []SectionalNeedleValue

func NewCompactSection(start NeedleId) *CompactSection {
	return &CompactSection{
		values:   make([]SectionalNeedleValue, batch),
		overflow: Overflow(make([]SectionalNeedleValue, 0)),
		start:    start,
	}
}

//return old entry size
func (cs *CompactSection) Set(key NeedleId, offset Offset, size uint32) (oldOffset Offset, oldSize uint32) {
	cs.Lock()
	if key > cs.end {
		cs.end = key
	}
	skey := SectionalNeedleId(key - cs.start)
	if i := cs.binarySearchValues(skey); i >= 0 {
		oldOffset, oldSize = cs.values[i].Offset, cs.values[i].Size
		//println("key", key, "old size", ret)
		cs.values[i].Offset, cs.values[i].Size = offset, size
	} else {
		needOverflow := cs.counter >= batch
		needOverflow = needOverflow || cs.counter > 0 && cs.values[cs.counter-1].Key > skey
		if needOverflow {
			//println("start", cs.start, "counter", cs.counter, "key", key)
			if oldValue, found := cs.overflow.findOverflowEntry(skey); found {
				oldOffset, oldSize = oldValue.Offset, oldValue.Size
			}
			cs.overflow = cs.overflow.setOverflowEntry(SectionalNeedleValue{Key: skey, Offset: offset, Size: size})
		} else {
			p := &cs.values[cs.counter]
			p.Key, p.Offset, p.Size = skey, offset, size
			//println("added index", cs.counter, "key", key, cs.values[cs.counter].Key)
			cs.counter++
		}
	}
	cs.Unlock()
	return
}

//return old entry size
func (cs *CompactSection) Delete(key NeedleId) uint32 {
	skey := SectionalNeedleId(key - cs.start)
	cs.Lock()
	ret := uint32(0)
	if i := cs.binarySearchValues(skey); i >= 0 {
		if cs.values[i].Size > 0 {
			ret = cs.values[i].Size
			cs.values[i].Size = 0
		}
	}
	if v, found := cs.overflow.findOverflowEntry(skey); found {
		cs.overflow = cs.overflow.deleteOverflowEntry(skey)
		ret = v.Size
	}
	cs.Unlock()
	return ret
}
func (cs *CompactSection) Get(key NeedleId) (*NeedleValue, bool) {
	cs.RLock()
	skey := SectionalNeedleId(key - cs.start)
	if v, ok := cs.overflow.findOverflowEntry(skey); ok {
		cs.RUnlock()
		nv := v.toNeedleValue(cs)
		return &nv, true
	}
	if i := cs.binarySearchValues(skey); i >= 0 {
		cs.RUnlock()
		nv := cs.values[i].toNeedleValue(cs)
		return &nv, true
	}
	cs.RUnlock()
	return nil, false
}
func (cs *CompactSection) binarySearchValues(key SectionalNeedleId) int {
	x := sort.Search(cs.counter, func(i int) bool {
		return cs.values[i].Key >= key
	})
	if x == cs.counter {
		return -1
	}
	if cs.values[x].Key > key {
		return -2
	}
	return x
}

//This map assumes mostly inserting increasing keys
//This map assumes mostly inserting increasing keys
type CompactMap struct {
	list []*CompactSection
}

func NewCompactMap() *CompactMap {
	return &CompactMap{}
}

func (cm *CompactMap) Set(key NeedleId, offset Offset, size uint32) (oldOffset Offset, oldSize uint32) {
	x := cm.binarySearchCompactSection(key)
	if x < 0 || (key-cm.list[x].start) > SectionalNeedleIdLimit {
		// println(x, "adding to existing", len(cm.list), "sections, starting", key)
		cs := NewCompactSection(key)
		cm.list = append(cm.list, cs)
		x = len(cm.list) - 1
		//keep compact section sorted by start
		for x >= 0 {
			if x > 0 && cm.list[x-1].start > key {
				cm.list[x] = cm.list[x-1]
				// println("shift", x, "start", cs.start, "to", x-1)
				x = x - 1
			} else {
				cm.list[x] = cs
				// println("cs", x, "start", cs.start)
				break
			}
		}
	}
	// println(key, "set to section[", x, "].start", cm.list[x].start)
	return cm.list[x].Set(key, offset, size)
}
func (cm *CompactMap) Delete(key NeedleId) uint32 {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		return uint32(0)
	}
	return cm.list[x].Delete(key)
}
func (cm *CompactMap) Get(key NeedleId) (*NeedleValue, bool) {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		return nil, false
	}
	return cm.list[x].Get(key)
}
func (cm *CompactMap) binarySearchCompactSection(key NeedleId) int {
	l, h := 0, len(cm.list)-1
	if h < 0 {
		return -5
	}
	if cm.list[h].start <= key {
		if cm.list[h].counter < batch || key <= cm.list[h].end {
			return h
		}
		return -4
	}
	for l <= h {
		m := (l + h) / 2
		if key < cm.list[m].start {
			h = m - 1
		} else { // cm.list[m].start <= key
			if cm.list[m+1].start <= key {
				l = m + 1
			} else {
				return m
			}
		}
	}
	return -3
}

// Visit visits all entries or stop if any error when visiting
func (cm *CompactMap) Visit(visit func(NeedleValue) error) error {
	for _, cs := range cm.list {
		cs.RLock()
		for _, v := range cs.overflow {
			if err := visit(v.toNeedleValue(cs)); err != nil {
				cs.RUnlock()
				return err
			}
		}
		for i, v := range cs.values {
			if i >= cs.counter {
				break
			}
			if _, found := cs.overflow.findOverflowEntry(v.Key); !found {
				if err := visit(v.toNeedleValue(cs)); err != nil {
					cs.RUnlock()
					return err
				}
			}
		}
		cs.RUnlock()
	}
	return nil
}

func (o Overflow) deleteOverflowEntry(key SectionalNeedleId) Overflow {
	length := len(o)
	deleteCandidate := sort.Search(length, func(i int) bool {
		return o[i].Key >= key
	})
	if deleteCandidate != length && o[deleteCandidate].Key == key {
		for i := deleteCandidate; i < length-1; i++ {
			o[i] = o[i+1]
		}
		o = o[0 : length-1]
	}
	return o
}

func (o Overflow) setOverflowEntry(needleValue SectionalNeedleValue) Overflow {
	insertCandidate := sort.Search(len(o), func(i int) bool {
		return o[i].Key >= needleValue.Key
	})
	if insertCandidate != len(o) && o[insertCandidate].Key == needleValue.Key {
		o[insertCandidate] = needleValue
	} else {
		o = append(o, needleValue)
		for i := len(o) - 1; i > insertCandidate; i-- {
			o[i] = o[i-1]
		}
		o[insertCandidate] = needleValue
	}
	return o
}

func (o Overflow) findOverflowEntry(key SectionalNeedleId) (nv SectionalNeedleValue, found bool) {
	foundCandidate := sort.Search(len(o), func(i int) bool {
		return o[i].Key >= key
	})
	if foundCandidate != len(o) && o[foundCandidate].Key == key {
		return o[foundCandidate], true
	}
	return nv, false
}

func (snv SectionalNeedleValue) toNeedleValue(cs *CompactSection) NeedleValue {
	return NeedleValue{NeedleId(snv.Key) + cs.start, snv.Offset, snv.Size}
}

func (nv NeedleValue) toSectionalNeedleValue(cs *CompactSection) SectionalNeedleValue {
	return SectionalNeedleValue{SectionalNeedleId(nv.Key - cs.start), nv.Offset, nv.Size}
}
<commit_msg>memory needle map mark size to be TombstoneFileSize<commit_after>package needle

import (
	. "github.com/chrislusf/seaweedfs/weed/storage/types"
	"sort"
	"sync"
)

const (
	batch = 100000
)

type SectionalNeedleId uint32

const SectionalNeedleIdLimit = 1<<32 - 1

type SectionalNeedleValue struct {
	Key    SectionalNeedleId
	Offset Offset `comment:"Volume offset"` //since aligned to 8 bytes, range is 4G*8=32G
	Size   uint32 `comment:"Size of the data portion"`
}

type CompactSection struct {
	sync.RWMutex
	values   []SectionalNeedleValue
	overflow Overflow
	start    NeedleId
	end      NeedleId
	counter  int
}

type Overflow []SectionalNeedleValue

func NewCompactSection(start NeedleId) *CompactSection {
	return &CompactSection{
		values:   make([]SectionalNeedleValue, batch),
		overflow: Overflow(make([]SectionalNeedleValue, 0)),
		start:    start,
	}
}

//return old entry size
func (cs *CompactSection) Set(key NeedleId, offset Offset, size uint32) (oldOffset Offset, oldSize uint32) {
	cs.Lock()
	if key > cs.end {
		cs.end = key
	}
	skey := SectionalNeedleId(key - cs.start)
	if i := cs.binarySearchValues(skey); i >= 0 {
		oldOffset, oldSize = cs.values[i].Offset, cs.values[i].Size
		//println("key", key, "old size", ret)
		cs.values[i].Offset, cs.values[i].Size = offset, size
	} else {
		needOverflow := cs.counter >= batch
		needOverflow = needOverflow || cs.counter > 0 && cs.values[cs.counter-1].Key > skey
		if needOverflow {
			//println("start", cs.start, "counter", cs.counter, "key", key)
			if oldValue, found := cs.overflow.findOverflowEntry(skey); found {
				oldOffset, oldSize = oldValue.Offset, oldValue.Size
			}
			cs.overflow = cs.overflow.setOverflowEntry(SectionalNeedleValue{Key: skey, Offset: offset, Size: size})
		} else {
			p := &cs.values[cs.counter]
			p.Key, p.Offset, p.Size = skey, offset, size
			//println("added index", cs.counter, "key", key, cs.values[cs.counter].Key)
			cs.counter++
		}
	}
	cs.Unlock()
	return
}

//return old entry size
func (cs *CompactSection) Delete(key NeedleId) uint32 {
	skey := SectionalNeedleId(key - cs.start)
	cs.Lock()
	ret := uint32(0)
	if i := cs.binarySearchValues(skey); i >= 0 {
		if cs.values[i].Size > 0 && cs.values[i].Size != TombstoneFileSize {
			ret = cs.values[i].Size
			cs.values[i].Size = TombstoneFileSize
		}
	}
	if v, found := cs.overflow.findOverflowEntry(skey); found {
		cs.overflow = cs.overflow.deleteOverflowEntry(skey)
		ret = v.Size
	}
	cs.Unlock()
	return ret
}
func (cs *CompactSection) Get(key NeedleId) (*NeedleValue, bool) {
	cs.RLock()
	skey := SectionalNeedleId(key - cs.start)
	if v, ok := cs.overflow.findOverflowEntry(skey); ok {
		cs.RUnlock()
		nv := v.toNeedleValue(cs)
		return &nv, true
	}
	if i := cs.binarySearchValues(skey); i >= 0 {
		cs.RUnlock()
		nv := cs.values[i].toNeedleValue(cs)
		return &nv, true
	}
	cs.RUnlock()
	return nil, false
}
func (cs *CompactSection) binarySearchValues(key SectionalNeedleId) int {
	x := sort.Search(cs.counter, func(i int) bool {
		return cs.values[i].Key >= key
	})
	if x == cs.counter {
		return -1
	}
	if cs.values[x].Key > key {
		return -2
	}
	return x
}

//This map assumes mostly inserting increasing keys
//This map assumes mostly inserting increasing keys
type CompactMap struct {
	list []*CompactSection
}

func NewCompactMap() *CompactMap {
	return &CompactMap{}
}

func (cm *CompactMap) Set(key NeedleId, offset Offset, size uint32) (oldOffset Offset, oldSize uint32) {
	x := cm.binarySearchCompactSection(key)
	if x < 0 || (key-cm.list[x].start) > SectionalNeedleIdLimit {
		// println(x, "adding to existing", len(cm.list), "sections, starting", key)
		cs := NewCompactSection(key)
		cm.list = append(cm.list, cs)
		x = len(cm.list) - 1
		//keep compact section sorted by start
		for x >= 0 {
			if x > 0 && cm.list[x-1].start > key {
				cm.list[x] = cm.list[x-1]
				// println("shift", x, "start", cs.start, "to", x-1)
				x = x - 1
			} else {
				cm.list[x] = cs
				// println("cs", x, "start", cs.start)
				break
			}
		}
	}
	// println(key, "set to section[", x, "].start", cm.list[x].start)
	return cm.list[x].Set(key, offset, size)
}
func (cm *CompactMap) Delete(key NeedleId) uint32 {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		return uint32(0)
	}
	return cm.list[x].Delete(key)
}
func (cm *CompactMap) Get(key NeedleId) (*NeedleValue, bool) {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		return nil, false
	}
	return cm.list[x].Get(key)
}
func (cm *CompactMap) binarySearchCompactSection(key NeedleId) int {
	l, h := 0, len(cm.list)-1
	if h < 0 {
		return -5
	}
	if cm.list[h].start <= key {
		if cm.list[h].counter < batch || key <= cm.list[h].end {
			return h
		}
		return -4
	}
	for l <= h {
		m := (l + h) / 2
		if key < cm.list[m].start {
			h = m - 1
		} else { // cm.list[m].start <= key
			if cm.list[m+1].start <= key {
				l = m + 1
			} else {
				return m
			}
		}
	}
	return -3
}

// Visit visits all entries or stop if any error when visiting
func (cm *CompactMap) Visit(visit func(NeedleValue) error) error {
	for _, cs := range cm.list {
		cs.RLock()
		for _, v := range cs.overflow {
			if err := visit(v.toNeedleValue(cs)); err != nil {
				cs.RUnlock()
				return err
			}
		}
		for i, v := range cs.values {
			if i >= cs.counter {
				break
			}
			if _, found := cs.overflow.findOverflowEntry(v.Key); !found {
				if err := visit(v.toNeedleValue(cs)); err != nil {
					cs.RUnlock()
					return err
				}
			}
		}
		cs.RUnlock()
	}
	return nil
}

func (o Overflow) deleteOverflowEntry(key SectionalNeedleId) Overflow {
	length := len(o)
	deleteCandidate := sort.Search(length, func(i int) bool {
		return o[i].Key >= key
	})
	if deleteCandidate != length && o[deleteCandidate].Key == key {
		for i := deleteCandidate; i < length-1; i++ {
			o[i] = o[i+1]
		}
		o = o[0 : length-1]
	}
	return o
}

func (o Overflow) setOverflowEntry(needleValue SectionalNeedleValue) Overflow {
	insertCandidate := sort.Search(len(o), func(i int) bool {
		return o[i].Key >= needleValue.Key
	})
	if insertCandidate != len(o) && o[insertCandidate].Key == needleValue.Key {
		o[insertCandidate] = needleValue
	} else {
		o = append(o, needleValue)
		for i := len(o) - 1; i > insertCandidate; i-- {
			o[i] = o[i-1]
		}
		o[insertCandidate] = needleValue
	}
	return o
}

func (o Overflow) findOverflowEntry(key SectionalNeedleId) (nv SectionalNeedleValue, found bool) {
	foundCandidate := sort.Search(len(o), func(i int) bool {
		return o[i].Key >= key
	})
	if foundCandidate != len(o) && o[foundCandidate].Key == key {
		return o[foundCandidate], true
	}
	return nv, false
}

func (snv SectionalNeedleValue) toNeedleValue(cs *CompactSection) NeedleValue {
	return NeedleValue{NeedleId(snv.Key) + cs.start, snv.Offset, snv.Size}
}

func (nv NeedleValue) toSectionalNeedleValue(cs *CompactSection) SectionalNeedleValue {
	return SectionalNeedleValue{SectionalNeedleId(nv.Key - cs.start), nv.Offset, nv.Size}
}
<|endoftext|>


<commit_before>package tempCrit

import (
	"math"
)
import (
	"../fit"
	"../solve"
	"../tempAll"
	vec "../vector"
)

type omegaFunc func(*tempAll.Environment, vec.Vector) (float64, error)

// Calculate the coefficients in the small-q pair dispersion relation:
//
// omega(q) = ax*q_x^2 + ay*q_y^2 + b*q_z^2 - mu_b
//
// The returned vector has the values {ax, ay, b, mu_b}. Due to x<->y symmetry
// we expect ax == ay.
func OmegaFit(env *tempAll.Environment, fn omegaFunc) (vec.Vector, error) {
	points := omegaCoeffsPoints()
	// evaluate omega_+(k) at each point
	omegas := make([]float64, len(points))
	X := make([]vec.Vector, len(points))
	for i, q := range points {
		var err error
		omegas[i], err = OmegaPlus(env, q)
		if err != nil {
			return nil, err
		}
		X[i] = vec.ZeroVector(4)
		X[i][0] = q[0] * q[0]
		X[i][1] = q[1] * q[1]
		X[i][2] = q[2] * q[2]
		X[i][3] = -1
	}
	return fit.Linear(omegas, X), nil
}

// Return a list of all k points surveyed by OmegaCoeffs().
func omegaCoeffsPoints() []vec.Vector {
	sk := 0.01 // small value of k
	ssk := sk / math.Sqrt(2)
	// unique point
	zero := vec.ZeroVector(3)
	// basis vectors
	xb := []float64{sk, 0.0, 0.0}
	yb := []float64{0.0, sk, 0.0}
	zb := []float64{0.0, 0.0, sk}
	xyb := []float64{ssk, ssk, 0.0}
	xzb := []float64{ssk, 0.0, ssk}
	yzb := []float64{0.0, ssk, ssk}
	basis := []vec.Vector{xb, yb, zb, xyb, xzb, yzb}
	// create points from basis
	numRadialPoints := 3
	points := []vec.Vector{zero}
	for _, v := range basis {
		for i := 1; i <= numRadialPoints; i++ {
			points = append(points, v.Mul(float64(i)))
		}
	}
	return points
}

// Calculate omega_+(k) by finding zeros of 1 - lambda_+
func OmegaPlus(env *tempAll.Environment, k vec.Vector) (float64, error) {
	lp := lambdaPlusFn(env, k)
	return solve.OneDimDiffRoot(lp, 0.01, 1e-9, 1e-9)
}

// Calculate omega_-(k) by finding zeros of 1 - lambda_-
func OmegaMinus(env *tempAll.Environment, k vec.Vector) (float64, error) {
	lm := lambdaMinusFn(env, k)
	return solve.OneDimDiffRoot(lm, 0.01, 1e-9, 1e-9)
}

// Create a function which calculates 1 - lambda_+(k, omega) with fixed k
func lambdaPlusFn(env *tempAll.Environment, k vec.Vector) func(float64) (float64, error) {
	return func(omega float64) (float64, error) {
		u, v := lambdaParts(env, k, omega)
		return 1.0 - (u + v), nil
	}
}

// Create a function which calculates 1 - lambda_-(k, omega) with fixed k
func lambdaMinusFn(env *tempAll.Environment, k vec.Vector) func(float64) (float64, error) {
	return func(omega float64) (float64, error) {
		u, v := lambdaParts(env, k, omega)
		return 1.0 - (u - v), nil
	}
}

// Calculate u, v in lambda_+/- = u +/- v
func lambdaParts(env *tempAll.Environment, k vec.Vector, omega float64) (float64, float64) {
	cx, cy, cz := math.Cos(k[0]), math.Cos(k[1]), math.Cos(k[2])
	Ex := 2.0 * (env.T0*cy + env.Tz*cz)
	Ey := 2.0 * (env.T0*cx + env.Tz*cz)
	Pis := Pi(env, []float64{k[0], k[1]}, omega)
	u := 0.5 * (Ex*Pis[0] + Ey*Pis[2])
	v := math.Sqrt(0.25*math.Pow(Ex*Pis[0]-Ey*Pis[2], 2.0) + Ex*Ey*math.Pow(Pis[1], 2.0))
	return u, v
}
<commit_msg>export OmegaFunc type<commit_after>package tempCrit

import (
	"math"
)
import (
	"../fit"
	"../solve"
	"../tempAll"
	vec "../vector"
)

type OmegaFunc func(*tempAll.Environment, vec.Vector) (float64, error)

// Calculate the coefficients in the small-q pair dispersion relation:
//
// omega(q) = ax*q_x^2 + ay*q_y^2 + b*q_z^2 - mu_b
//
// The returned vector has the values {ax, ay, b, mu_b}. Due to x<->y symmetry
// we expect ax == ay.
func OmegaFit(env *tempAll.Environment, fn OmegaFunc) (vec.Vector, error) {
	points := omegaCoeffsPoints()
	// evaluate omega_+(k) at each point
	omegas := make([]float64, len(points))
	X := make([]vec.Vector, len(points))
	for i, q := range points {
		var err error
		omegas[i], err = OmegaPlus(env, q)
		if err != nil {
			return nil, err
		}
		X[i] = vec.ZeroVector(4)
		X[i][0] = q[0] * q[0]
		X[i][1] = q[1] * q[1]
		X[i][2] = q[2] * q[2]
		X[i][3] = -1
	}
	return fit.Linear(omegas, X), nil
}

// Return a list of all k points surveyed by OmegaCoeffs().
func omegaCoeffsPoints() []vec.Vector {
	sk := 0.01 // small value of k
	ssk := sk / math.Sqrt(2)
	// unique point
	zero := vec.ZeroVector(3)
	// basis vectors
	xb := []float64{sk, 0.0, 0.0}
	yb := []float64{0.0, sk, 0.0}
	zb := []float64{0.0, 0.0, sk}
	xyb := []float64{ssk, ssk, 0.0}
	xzb := []float64{ssk, 0.0, ssk}
	yzb := []float64{0.0, ssk, ssk}
	basis := []vec.Vector{xb, yb, zb, xyb, xzb, yzb}
	// create points from basis
	numRadialPoints := 3
	points := []vec.Vector{zero}
	for _, v := range basis {
		for i := 1; i <= numRadialPoints; i++ {
			points = append(points, v.Mul(float64(i)))
		}
	}
	return points
}

// Calculate omega_+(k) by finding zeros of 1 - lambda_+
func OmegaPlus(env *tempAll.Environment, k vec.Vector) (float64, error) {
	lp := lambdaPlusFn(env, k)
	return solve.OneDimDiffRoot(lp, 0.01, 1e-9, 1e-9)
}

// Calculate omega_-(k) by finding zeros of 1 - lambda_-
func OmegaMinus(env *tempAll.Environment, k vec.Vector) (float64, error) {
	lm := lambdaMinusFn(env, k)
	return solve.OneDimDiffRoot(lm, 0.01, 1e-9, 1e-9)
}

// Create a function which calculates 1 - lambda_+(k, omega) with fixed k
func lambdaPlusFn(env *tempAll.Environment, k vec.Vector) func(float64) (float64, error) {
	return func(omega float64) (float64, error) {
		u, v := lambdaParts(env, k, omega)
		return 1.0 - (u + v), nil
	}
}

// Create a function which calculates 1 - lambda_-(k, omega) with fixed k
func lambdaMinusFn(env *tempAll.Environment, k vec.Vector) func(float64) (float64, error) {
	return func(omega float64) (float64, error) {
		u, v := lambdaParts(env, k, omega)
		return 1.0 - (u - v), nil
	}
}

// Calculate u, v in lambda_+/- = u +/- v
func lambdaParts(env *tempAll.Environment, k vec.Vector, omega float64) (float64, float64) {
	cx, cy, cz := math.Cos(k[0]), math.Cos(k[1]), math.Cos(k[2])
	Ex := 2.0 * (env.T0*cy + env.Tz*cz)
	Ey := 2.0 * (env.T0*cx + env.Tz*cz)
	Pis := Pi(env, []float64{k[0], k[1]}, omega)
	u := 0.5 * (Ex*Pis[0] + Ey*Pis[2])
	v := math.Sqrt(0.25*math.Pow(Ex*Pis[0]-Ey*Pis[2], 2.0) + Ex*Ey*math.Pow(Pis[1], 2.0))
	return u, v
}
<|endoftext|>


<commit_before>package maxd

import "github.com/catorpilor/leetcode/utils"

func maxDepth(root *utils.TreeNode) int {
	return 0
}
<commit_msg>solve 104 using recursive dfs<commit_after>package maxd

import "github.com/catorpilor/leetcode/utils"

func maxDepth(root *utils.TreeNode) int {
	return useDfs(root)
}

// useDfs time complexity O(N), space complexity O(lgN)
func useDfs(root *utils.TreeNode) int {
	if root == nil {
		return 0
	}
	l := useDfs(root.Left)
	r := useDfs(root.Right)
	return 1 + utils.Max(l, r)
}
<|endoftext|>