<commit_before>/*
Copyright 2019 Google LLC

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

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

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

package operations

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"strings"

	"go.uber.org/zap"
	"knative.dev/pkg/kmeta"
	"knative.dev/pkg/logging"

	storageClient "cloud.google.com/go/storage"
	"github.com/google/knative-gcp/pkg/operations"
	"google.golang.org/grpc/codes"
	gstatus "google.golang.org/grpc/status"

	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var (
	// Mapping of the storage importer eventTypes to google storage types.
	storageEventTypes = map[string]string{
		"finalize":       "OBJECT_FINALIZE",
		"archive":        "OBJECT_ARCHIVE",
		"delete":         "OBJECT_DELETE",
		"metadataUpdate": "OBJECT_METADATA_UPDATE",
	}
)

type NotificationActionResult struct {
	// Result is the result the operation attempted.
	Result bool `json:"result,omitempty"`
	// Error is the error string if failure occurred
	Error string `json:"error,omitempty"`
	// NotificationId holds the notification ID for GCS
	// and is filled in during create operation.
	NotificationId string `json:"notificationId,omitempty"`
}

// NotificationArgs are the configuration required to make a NewNotificationOps.
type NotificationArgs struct {
	// UID of the resource that caused the action to be taken. Will
	// be added as a label to the podtemplate.
	UID string
	// Image is the actual binary that we'll run to operate on the
	// notification.
	Image string
	// Action is what the binary should do
	Action    string
	ProjectID string
	// Bucket
	Bucket string
	// Topic we'll use for pubsub target.
	TopicID string
	// NotificationId is the notifification ID that GCS gives
	// back to us. We need that to delete it.
	NotificationId string
	// EventTypes is an array of strings specifying which
	// event types we want the notification to fire on.
	EventTypes []string
	// ObjectNamePrefix is an optional filter
	ObjectNamePrefix string
	// CustomAttributes is the list of additional attributes to have
	// GCS supply back to us when it sends a notification.
	CustomAttributes map[string]string
	Secret           corev1.SecretKeySelector
	Owner            kmeta.OwnerRefable
}

// NewNotificationOps returns a new batch Job resource.
func NewNotificationOps(arg NotificationArgs) *batchv1.Job {
	env := []corev1.EnvVar{{
		Name:  "ACTION",
		Value: arg.Action,
	}, {
		Name:  "PROJECT_ID",
		Value: arg.ProjectID,
	}, {
		Name:  "BUCKET",
		Value: arg.Bucket,
	}, {
		Name:  "PUBSUB_TOPIC_ID",
		Value: arg.TopicID,
	}}

	switch arg.Action {
	case operations.ActionCreate:
		env = append(env, []corev1.EnvVar{
			{
				Name:  "EVENT_TYPES",
				Value: strings.Join(arg.EventTypes, ":"),
			}, {
				Name:  "PUBSUB_TOPIC_ID",
				Value: arg.TopicID,
			}}...)
	case operations.ActionDelete:
		env = append(env, []corev1.EnvVar{{
			Name:  "NOTIFICATION_ID",
			Value: arg.NotificationId,
		}}...)
	}

	podTemplate := operations.MakePodTemplate(arg.Image, arg.UID, arg.Action, arg.Secret, env...)

	backoffLimit := int32(3)
	parallelism := int32(1)

	return &batchv1.Job{
		ObjectMeta: metav1.ObjectMeta{
			Name:            NotificationJobName(arg.Owner, arg.Action),
			Namespace:       arg.Owner.GetObjectMeta().GetNamespace(),
			Labels:          NotificationJobLabels(arg.Owner, arg.Action),
			OwnerReferences: []metav1.OwnerReference{*kmeta.NewControllerRef(arg.Owner)},
		},
		Spec: batchv1.JobSpec{
			BackoffLimit: &backoffLimit,
			Parallelism:  &parallelism,
			Template:     *podTemplate,
		},
	}
}

// TODO: the job could output the resolved projectID.

// NotificationOps defines the configuration to use for this operation.
type NotificationOps struct {
	StorageOps

	// Action is the operation the job should run.
	// Options: [exists, create, delete]
	Action string `envconfig:"ACTION" required:"true"`

	// Topic is the environment variable containing the PubSub Topic being
	// subscribed to's name. In the form that is unique within the project.
	// E.g. 'laconia', not 'projects/my-gcp-project/topics/laconia'.
	Topic string `envconfig:"PUBSUB_TOPIC_ID" required:"true"`

	// Bucket to operate on
	Bucket string `envconfig:"BUCKET" required:"true"`

	// NotificationId is the environment variable containing the name of the
	// subscription to use.
	NotificationId string `envconfig:"NOTIFICATION_ID" required:"false" default:""`

	// EventTypes is a : separated eventtypes, if omitted all will be used.
	// TODO: Look at native envconfig list support
	EventTypes string `envconfig:"EVENT_TYPES" required:"false" default:""`

	// ObjectNamePrefix is an optional filter for the GCS
	ObjectNamePrefix string `envconfig:"OBJECT_NAME_PREFIX" required:"false" default:""`

	// TODO; Add support for custom attributes. Look at using envconfig Map with
	// necessary encoding / decoding.
}

//var (
//	ignoreNotificationConfig = cmpopts.IgnoreFields(pubsub.NotificationConfig{}, "Topic", "Labels")
//)

// Run will perform the action configured upon a subscription.
func (n *NotificationOps) Run(ctx context.Context) error {
	if n.Client == nil {
		return errors.New("pub/sub client is nil")
	}
	logger := logging.FromContext(ctx)

	logger = logger.With(
		zap.String("action", n.Action),
		zap.String("project", n.Project),
		zap.String("topic", n.Topic),
		zap.String("subscription", n.NotificationId),
	)

	logger.Info("Storage Notification Job.")

	// Load the Bucket.
	bucket := n.Client.Bucket(n.Bucket)

	switch n.Action {
	case operations.ActionExists:
		// If notification doesn't exist, that is an error.
		logger.Info("Previously created.")

	case operations.ActionCreate:
		customAttributes := make(map[string]string)
		// for k, v := range n.CustomAttributes {
		//			customAttributes[k] = v
		//		}

		// Add our own event type here...
		customAttributes["knative-gcp"] = "google.storage"

		eventTypes := strings.Split(n.EventTypes, ":")
		logger.Infof("Creating a notification on bucket %s", n.Bucket)

		nc := storageClient.Notification{
			TopicProjectID:   n.Project,
			TopicID:          n.Topic,
			PayloadFormat:    storageClient.JSONPayload,
			EventTypes:       n.toStorageEventTypes(eventTypes),
			ObjectNamePrefix: n.ObjectNamePrefix,
			CustomAttributes: customAttributes,
		}

		logger.Info("NOTIFIcATION IS: %+v", nc)
		notification, err := bucket.AddNotification(ctx, &nc)
		if err != nil {
			result := &NotificationActionResult{
				Result: false,
				Error:  err.Error(),
			}
			logger.Infof("Failed to create Notification: %s", err)
			err = n.writeTerminationMessage(result)
			return err
		}
		logger.Infof("Created Notification %q", notification.ID)
		result := &NotificationActionResult{
			Result:         true,
			NotificationId: notification.ID,
		}
		err = n.writeTerminationMessage(result)
		if err != nil {
			logger.Infof("Failed to write termination message: %s", err)
			return err
		}
		/*
			// If topic doesn't exist, create it.
			if !exists {
				// Create a new subscription to the previous topic with the given name.
				sub, err = s.Client.CreateNotification(ctx, s.Notification, subConfig)
				if err != nil {
					return fmt.Errorf("failed to create subscription, %s", err)
				}
				logger.Info("Successfully created.")
			} else {
				// TODO: here is where we could update config.
				logger.Info("Previously created.")
				// Get current config.
				currentConfig, err := sub.Config(ctx)
				if err != nil {
					return fmt.Errorf("failed to get subscription config, %s", err)
				}
				// Compare the current config to the expected config. Update if different.
				if diff := cmp.Diff(subConfig, currentConfig, ignoreSubConfig); diff != "" {
					_, err := sub.Update(ctx, pubsub.NotificationConfig{
						AckDeadline:         s.AckDeadline,
						RetainAckedMessages: s.RetainAckedMessages,
						RetentionDuration:   s.RetentionDuration,
						Labels:              currentConfig.Labels,
					})
					if err != nil {
						return fmt.Errorf("failed to update subscription config, %s", err)
					}
					logger.Info("Updated subscription config.", zap.String("diff", diff))

				}
			}
		*/

	case operations.ActionDelete:
		notifications, err := bucket.Notifications(ctx)
		if err != nil {
			logger.Infof("Failed to fetch existing notifications: %s", err)
			return err
		}

		// This is bit wonky because, we could always just try to delete, but figuring out
		// if an error returned is NotFound seems to not really work, so, we'll try
		// checking first the list and only then deleting.
		notificationId := n.NotificationId
		if notificationId != "" {
			if existing, ok := notifications[notificationId]; ok {
				logger.Infof("Found existing notification: %+v", existing)
				logger.Infof("Deleting notification as: %q", notificationId)
				err = bucket.DeleteNotification(ctx, notificationId)
				if err == nil {
					logger.Infof("Deleted Notification: %q", notificationId)
					err = n.writeTerminationMessage(&NotificationActionResult{Result: true})
					if err != nil {
						logger.Infof("Failed to write termination message: %s", err)
						return err
					}
					return nil
				}

				if st, ok := gstatus.FromError(err); !ok {
					logger.Infof("error from the cloud storage client: %s", err)
					writeErr := n.writeTerminationMessage(&NotificationActionResult{Result: false, Error: err.Error()})
					if writeErr != nil {
						logger.Infof("Failed to write termination message: %s", writeErr)
						return err
					}
					return err
				} else if st.Code() != codes.NotFound {
					writeErr := n.writeTerminationMessage(&NotificationActionResult{Result: false, Error: err.Error()})
					if writeErr != nil {
						logger.Infof("Failed to write termination message: %s", writeErr)
						return err
					}
					return err
				}
			}
		}
	default:
		return fmt.Errorf("unknown action value %v", n.Action)
	}

	logger.Info("Done.")
	return nil
}

func (n *NotificationOps) toStorageEventTypes(eventTypes []string) []string {
	storageTypes := make([]string, 0, len(eventTypes))
	for _, eventType := range eventTypes {
		storageTypes = append(storageTypes, storageEventTypes[eventType])
	}

	if len(storageTypes) == 0 {
		return append(storageTypes, "OBJECT_FINALIZE")
	}
	return storageTypes
}

//func
func (n *NotificationOps) writeTerminationMessage(result *NotificationActionResult) error {
	m, err := json.Marshal(result)
	if err != nil {
		return err
	}
	return ioutil.WriteFile("/dev/termination-log", m, 0644)
}
<commit_msg>remove more cruft<commit_after>/*
Copyright 2019 Google LLC

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

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

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

package operations

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"strings"

	"go.uber.org/zap"
	"knative.dev/pkg/kmeta"
	"knative.dev/pkg/logging"

	storageClient "cloud.google.com/go/storage"
	"github.com/google/knative-gcp/pkg/operations"
	"google.golang.org/grpc/codes"
	gstatus "google.golang.org/grpc/status"

	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var (
	// Mapping of the storage importer eventTypes to google storage types.
	storageEventTypes = map[string]string{
		"finalize":       "OBJECT_FINALIZE",
		"archive":        "OBJECT_ARCHIVE",
		"delete":         "OBJECT_DELETE",
		"metadataUpdate": "OBJECT_METADATA_UPDATE",
	}
)

// TODO: the job could output the resolved projectID.
type NotificationActionResult struct {
	// Result is the result the operation attempted.
	Result bool `json:"result,omitempty"`
	// Error is the error string if failure occurred
	Error string `json:"error,omitempty"`
	// NotificationId holds the notification ID for GCS
	// and is filled in during create operation.
	NotificationId string `json:"notificationId,omitempty"`
}

// NotificationArgs are the configuration required to make a NewNotificationOps.
type NotificationArgs struct {
	// UID of the resource that caused the action to be taken. Will
	// be added as a label to the podtemplate.
	UID string
	// Image is the actual binary that we'll run to operate on the
	// notification.
	Image string
	// Action is what the binary should do
	Action    string
	ProjectID string
	// Bucket
	Bucket string
	// Topic we'll use for pubsub target.
	TopicID string
	// NotificationId is the notifification ID that GCS gives
	// back to us. We need that to delete it.
	NotificationId string
	// EventTypes is an array of strings specifying which
	// event types we want the notification to fire on.
	EventTypes []string
	// ObjectNamePrefix is an optional filter
	ObjectNamePrefix string
	// CustomAttributes is the list of additional attributes to have
	// GCS supply back to us when it sends a notification.
	CustomAttributes map[string]string
	Secret           corev1.SecretKeySelector
	Owner            kmeta.OwnerRefable
}

// NewNotificationOps returns a new batch Job resource.
func NewNotificationOps(arg NotificationArgs) *batchv1.Job {
	env := []corev1.EnvVar{{
		Name:  "ACTION",
		Value: arg.Action,
	}, {
		Name:  "PROJECT_ID",
		Value: arg.ProjectID,
	}, {
		Name:  "BUCKET",
		Value: arg.Bucket,
	}, {
		Name:  "PUBSUB_TOPIC_ID",
		Value: arg.TopicID,
	}}

	switch arg.Action {
	case operations.ActionCreate:
		env = append(env, []corev1.EnvVar{
			{
				Name:  "EVENT_TYPES",
				Value: strings.Join(arg.EventTypes, ":"),
			}, {
				Name:  "PUBSUB_TOPIC_ID",
				Value: arg.TopicID,
			}}...)
	case operations.ActionDelete:
		env = append(env, []corev1.EnvVar{{
			Name:  "NOTIFICATION_ID",
			Value: arg.NotificationId,
		}}...)
	}

	podTemplate := operations.MakePodTemplate(arg.Image, arg.UID, arg.Action, arg.Secret, env...)

	backoffLimit := int32(3)
	parallelism := int32(1)

	return &batchv1.Job{
		ObjectMeta: metav1.ObjectMeta{
			Name:            NotificationJobName(arg.Owner, arg.Action),
			Namespace:       arg.Owner.GetObjectMeta().GetNamespace(),
			Labels:          NotificationJobLabels(arg.Owner, arg.Action),
			OwnerReferences: []metav1.OwnerReference{*kmeta.NewControllerRef(arg.Owner)},
		},
		Spec: batchv1.JobSpec{
			BackoffLimit: &backoffLimit,
			Parallelism:  &parallelism,
			Template:     *podTemplate,
		},
	}
}

// NotificationOps defines the configuration to use for this operation.
type NotificationOps struct {
	StorageOps

	// Action is the operation the job should run.
	// Options: [exists, create, delete]
	Action string `envconfig:"ACTION" required:"true"`

	// Topic is the environment variable containing the PubSub Topic being
	// subscribed to's name. In the form that is unique within the project.
	// E.g. 'laconia', not 'projects/my-gcp-project/topics/laconia'.
	Topic string `envconfig:"PUBSUB_TOPIC_ID" required:"true"`

	// Bucket to operate on
	Bucket string `envconfig:"BUCKET" required:"true"`

	// NotificationId is the environment variable containing the name of the
	// subscription to use.
	NotificationId string `envconfig:"NOTIFICATION_ID" required:"false" default:""`

	// EventTypes is a : separated eventtypes, if omitted all will be used.
	// TODO: Look at native envconfig list support
	EventTypes string `envconfig:"EVENT_TYPES" required:"false" default:""`

	// ObjectNamePrefix is an optional filter for the GCS
	ObjectNamePrefix string `envconfig:"OBJECT_NAME_PREFIX" required:"false" default:""`

	// TODO; Add support for custom attributes. Look at using envconfig Map with
	// necessary encoding / decoding.
}

// Run will perform the action configured upon a subscription.
func (n *NotificationOps) Run(ctx context.Context) error {
	if n.Client == nil {
		return errors.New("pub/sub client is nil")
	}
	logger := logging.FromContext(ctx)

	logger = logger.With(
		zap.String("action", n.Action),
		zap.String("project", n.Project),
		zap.String("topic", n.Topic),
		zap.String("subscription", n.NotificationId),
	)

	logger.Info("Storage Notification Job.")

	// Load the Bucket.
	bucket := n.Client.Bucket(n.Bucket)

	switch n.Action {
	case operations.ActionExists:
		// If notification doesn't exist, that is an error.
		logger.Info("Previously created.")

	case operations.ActionCreate:
		customAttributes := make(map[string]string)
		// TODO: figure out how this works with envconfig esp. with encoding
		// values there.
		// for k, v := range n.CustomAttributes {
		//			customAttributes[k] = v
		//		}

		// Add our own event type here...
		customAttributes["knative-gcp"] = "google.storage"

		eventTypes := strings.Split(n.EventTypes, ":")
		logger.Infof("Creating a notification on bucket %s", n.Bucket)

		nc := storageClient.Notification{
			TopicProjectID:   n.Project,
			TopicID:          n.Topic,
			PayloadFormat:    storageClient.JSONPayload,
			EventTypes:       n.toStorageEventTypes(eventTypes),
			ObjectNamePrefix: n.ObjectNamePrefix,
			CustomAttributes: customAttributes,
		}

		logger.Info("NOTIFIcATION IS: %+v", nc)
		notification, err := bucket.AddNotification(ctx, &nc)
		if err != nil {
			result := &NotificationActionResult{
				Result: false,
				Error:  err.Error(),
			}
			logger.Infof("Failed to create Notification: %s", err)
			err = n.writeTerminationMessage(result)
			return err
		}
		logger.Infof("Created Notification %q", notification.ID)
		result := &NotificationActionResult{
			Result:         true,
			NotificationId: notification.ID,
		}
		err = n.writeTerminationMessage(result)
		if err != nil {
			logger.Infof("Failed to write termination message: %s", err)
			return err
		}
	case operations.ActionDelete:
		notifications, err := bucket.Notifications(ctx)
		if err != nil {
			logger.Infof("Failed to fetch existing notifications: %s", err)
			return err
		}

		// This is bit wonky because, we could always just try to delete, but figuring out
		// if an error returned is NotFound seems to not really work, so, we'll try
		// checking first the list and only then deleting.
		notificationId := n.NotificationId
		if notificationId != "" {
			if existing, ok := notifications[notificationId]; ok {
				logger.Infof("Found existing notification: %+v", existing)
				logger.Infof("Deleting notification as: %q", notificationId)
				err = bucket.DeleteNotification(ctx, notificationId)
				if err == nil {
					logger.Infof("Deleted Notification: %q", notificationId)
					err = n.writeTerminationMessage(&NotificationActionResult{Result: true})
					if err != nil {
						logger.Infof("Failed to write termination message: %s", err)
						return err
					}
					return nil
				}

				if st, ok := gstatus.FromError(err); !ok {
					logger.Infof("error from the cloud storage client: %s", err)
					writeErr := n.writeTerminationMessage(&NotificationActionResult{Result: false, Error: err.Error()})
					if writeErr != nil {
						logger.Infof("Failed to write termination message: %s", writeErr)
						return err
					}
					return err
				} else if st.Code() != codes.NotFound {
					writeErr := n.writeTerminationMessage(&NotificationActionResult{Result: false, Error: err.Error()})
					if writeErr != nil {
						logger.Infof("Failed to write termination message: %s", writeErr)
						return err
					}
					return err
				}
			}
		}
	default:
		return fmt.Errorf("unknown action value %v", n.Action)
	}

	logger.Info("Done.")
	return nil
}

func (n *NotificationOps) toStorageEventTypes(eventTypes []string) []string {
	storageTypes := make([]string, 0, len(eventTypes))
	for _, eventType := range eventTypes {
		storageTypes = append(storageTypes, storageEventTypes[eventType])
	}

	if len(storageTypes) == 0 {
		return append(storageTypes, "OBJECT_FINALIZE")
	}
	return storageTypes
}

func (n *NotificationOps) writeTerminationMessage(result *NotificationActionResult) error {
	m, err := json.Marshal(result)
	if err != nil {
		return err
	}
	return ioutil.WriteFile("/dev/termination-log", m, 0644)
}
<|endoftext|>


<commit_before>// MIT License
//
// Copyright (c) 2017 Ryan Fowler
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

// Package tree provides an implementation of a red-black tree.
package tree

// Int represents an integer that implements the Item interface.
type Int int

// Less returns true if the Int is less than the provided Int. If the provided
// Item is not an Int, Less will panic.
func (i Int) Less(than Item) bool {
	return i < than.(Int)
}

// Item is the interface that wraps the Less method.
//
// Less should return 'true' if the instance is "less than" the provided Item.
// Items are considered equal if neither are less than each other.
// E.g. Items 'a' & 'b' are considered equal if: (!a.Less(b) && !b.Less(a))
type Item interface {
	Less(Item) bool
}

// RedBlackTree is an in-memory implementation of a red-black tree.
//
// The internal data structure will automatically re-balance, and therefore
// allow for O(log(n)) retrieval, insertion, and deletion.
//
// Note: While read-only operations may occur concurrently, any write operation
// must be serially executed (typically protected with a mutex).
type RedBlackTree struct {
	root *node
	size int
}

// Ascend (O(n)) starts at the first Item and calls 'fn' for each Item until no
// Items remain or fn returns 'false'.
func (t *RedBlackTree) Ascend(fn func(Item) bool) {
	if t.root == nil {
		return
	}
	n := t.root.min()
	for n != nil && fn(n.item) {
		n = n.next()
	}
}

// Descend (O(n)) starts at the last Item and calls 'fn' for each Item until no
// Items remain or fn returns 'false'.
func (t *RedBlackTree) Descend(fn func(Item) bool) {
	if t.root == nil {
		return
	}
	n := t.root.max()
	for n != nil && fn(n.item) {
		n = n.prev()
	}
}

// Delete (O(log(n))) deletes an item in the RedBlackTree equal to the provided
// item. If an item was deleted, it is returned. Otherwise, nil is returned.
//
// Note: equality for items a & b is: (!a.Less(b) && !b.Less(a)).
func (t *RedBlackTree) Delete(item Item) Item {
	if t.root == nil {
		return nil
	}
	return t.root.deleteItem(t, item)
}

// DeleteMax (O(log(n))) deletes the maximum item in the RedBlackTree, returning
// it. If the tree is empty, nil is returned.
func (t *RedBlackTree) DeleteMax() Item {
	if t.root == nil {
		return nil
	}
	return t.root.deleteMax(t)
}

// DeleteMin (O(log(n))) deletes the minimum item in the RedBlackTree, returning
// it. If the tree is empty, nil is returned.
func (t *RedBlackTree) DeleteMin() Item {
	if t.root == nil {
		return nil
	}
	return t.root.deleteMin(t)
}

// Get (O(log(n))) retrieves an item in the RedBlackTree equal to the provided
// item. If an item was found, it is returned. Otherwise, nil is returned.
//
// Note: equality for items a & b is: (!a.Less(b) && !b.Less(a)).
func (t *RedBlackTree) Get(item Item) Item {
	n := t.root.find(item)
	if n == nil {
		return nil
	}
	return n.item
}

// Insert (O(log(n))) inserts (or replaces) an item into the RedBlackTree. If an
// item was replaced, it is returned. Otherwise, nil is returned.
//
// Note: equality for items a & b is: (!a.Less(b) && !b.Less(a)).
func (t *RedBlackTree) Insert(item Item) Item {
	if t.root == nil {
		t.root = newNode(nil, item)
		t.root.colour = colourBlack
		t.size++
		return nil
	}
	n, oldItem := t.root.insert(item)
	if oldItem == nil {
		t.size++
		n.rebalanceInsert(t)
	}
	return oldItem
}

// Exists (O(log(n))) returns 'true' if an item equal to the provided item
// exists in the RedBlackTree.
//
// Note: equality for items a & b is: (!a.Less(b) && !b.Less(a)).
func (t *RedBlackTree) Exists(item Item) bool {
	return t.Get(item) != nil
}

// Min (O(log(n))) returns the minimum item in the RedBlackTree. If the tree is
// empty, nil is returned.
func (t *RedBlackTree) Min() Item {
	if t.root == nil {
		return nil
	}
	n := t.root
	for n.left != nil {
		n = n.left
	}
	return n.item
}

// Max (O(log(n))) returns the maximum item in the RedBlackTree. If the tree is
// empty, nil is returned.
func (t *RedBlackTree) Max() Item {
	if t.root == nil {
		return nil
	}
	n := t.root
	for n.right != nil {
		n = n.right
	}
	return n.item
}

// Size (O(1)) returns the number of items in the RedBlackTree.
func (t *RedBlackTree) Size() int {
	return t.size
}

type colour uint8

const (
	colourRed   colour = 0
	colourBlack colour = 1
)

type node struct {
	colour      colour
	parent      *node
	left, right *node
	item        Item
}

func newNode(parent *node, item Item) *node {
	return &node{
		colour: colourRed,
		parent: parent,
		item:   item,
	}
}

func (n *node) find(item Item) *node {
	for n != nil {
		switch {
		case item.Less(n.item):
			n = n.left
		case n.item.Less(item):
			n = n.right
		default:
			return n
		}
	}
	return nil
}

func (n *node) deleteMax(t *RedBlackTree) Item {
	return n.max().deleteNode(t)
}

func (n *node) deleteMin(t *RedBlackTree) Item {
	return n.min().deleteNode(t)
}

func (n *node) deleteItem(t *RedBlackTree, item Item) Item {
	n = n.find(item)
	if n == nil {
		return nil
	}
	return n.deleteNode(t)
}

func (n *node) deleteNode(t *RedBlackTree) Item {
	t.size--
	delItem := n.item

	var child, parent *node
	for {
		if n.left == nil {
			child = n.right
			parent = n.parent
			n.replaceNode(t, n.right)
			break
		}
		if n.right == nil {
			child = n.left
			parent = n.parent
			n.replaceNode(t, n.left)
			break
		}
		// replace minimum value in right subtree with node to delete.
		min := n.right.min()
		n.item = min.item
		n = min
	}

	if n.isRed() {
		return delItem
	}
	if child.isRed() {
		child.colour = colourBlack
		return delItem
	}
	child.rebalanceDelete(t, parent)
	return delItem
}

func (n *node) rebalanceDelete(t *RedBlackTree, parent *node) {
	var s *node
	for {
		// Case 1.
		if n == t.root {
			return
		}
		if n != nil {
			parent = n.parent
		}
		// Case 2.
		s = n.sibling(parent)
		if s.isRed() {
			parent.colour = colourRed
			s.colour = colourBlack
			if n == parent.left {
				parent.rotateLeft(t)
			} else {
				parent.rotateRight(t)
			}
		}
		// Case 3.
		s = n.sibling(parent)
		if parent.isBlack() && s.isBlack() && s != nil && s.left.isBlack() && s.right.isBlack() {
			s.colour = colourRed
			n = parent
			if n != nil {
				parent = n.parent
			} else {
				parent = nil
			}
			continue
		}
		break
	}
	// Case 4.
	if parent.isRed() &&
		s.isBlack() &&
		s != nil &&
		s.left.isBlack() &&
		s.right.isBlack() {
		s.colour = colourRed
		parent.colour = colourBlack
		return
	}
	// Case 5.
	if s.isBlack() && s != nil {
		if n == parent.left && s.right.isBlack() && s.left.isRed() {
			s.colour = colourRed
			s.left.colour = colourBlack
			s.rotateRight(t)
		} else if n == parent.right && s.left.isBlack() && s.right.isRed() {
			s.colour = colourRed
			s.right.colour = colourBlack
			s.rotateLeft(t)
		}
	}
	// Case 6.
	s = n.sibling(parent)
	if s != nil {
		s.colour = parent.colour
		parent.colour = colourBlack
		if n == parent.left {
			s.right.colour = colourBlack
			parent.rotateLeft(t)
		} else {
			s.left.colour = colourBlack
			parent.rotateRight(t)
		}
	}
}

func (n *node) isRed() bool {
	return n != nil && n.colour == colourRed
}

func (n *node) isBlack() bool {
	return n == nil || n.colour == colourBlack
}

func (n *node) sibling(parent *node) *node {
	if n == parent.left {
		return parent.right
	}
	return parent.left
}

func (n *node) replaceNode(t *RedBlackTree, child *node) {
	switch {
	case n.parent == nil:
		t.root = child
	case n == n.parent.left:
		n.parent.left = child
	default:
		n.parent.right = child
	}
	if child != nil {
		child.parent = n.parent
	}
}

func (n *node) min() *node {
	for n.left != nil {
		n = n.left
	}
	return n
}

func (n *node) max() *node {
	for n.right != nil {
		n = n.right
	}
	return n
}

func (n *node) next() *node {
	if n.right != nil {
		return n.right.min()
	}
	parent := n.parent
	for parent != nil && parent.right == n {
		n = parent
		parent = n.parent
	}
	return parent
}

func (n *node) prev() *node {
	if n.left != nil {
		return n.left.max()
	}
	parent := n.parent
	for parent != nil && parent.left == n {
		n = parent
		parent = n.parent
	}
	return parent
}

func (n *node) insert(item Item) (*node, Item) {
	for {
		switch {
		case item.Less(n.item):
			if n.left == nil {
				n.left = newNode(n, item)
				return n.left, nil
			}
			n = n.left
		case n.item.Less(item):
			if n.right == nil {
				n.right = newNode(n, item)
				return n.right, nil
			}
			n = n.right
		default:
			oldItem := n.item
			n.item = item
			return n, oldItem
		}
	}
}

func (n *node) rebalanceInsert(t *RedBlackTree) {
	var g *node
	for {
		// Case 1.
		if n.parent == nil {
			n.colour = colourBlack
			return
		}
		// Case 2.
		if n.parent.colour == colourBlack {
			return
		}
		// Case 3.
		g = n.grandparent()
		var ps *node
		if g != nil {
			if n.parent == g.left {
				ps = g.right
			} else {
				ps = g.left
			}
		}
		if ps == nil || ps.colour == colourBlack {
			break
		}
		n.parent.colour = colourBlack
		ps.colour = colourBlack
		g.colour = colourRed
		n = g
	}
	// Case 4.
	if n == n.parent.right && n.parent == g.left {
		n.parent.rotateLeft(t)
		n = n.left
		g = n.grandparent()
	} else if n == n.parent.left && n.parent == g.right {
		n.parent.rotateRight(t)
		n = n.right
		g = n.grandparent()
	}
	// Case 5.
	n.parent.colour = colourBlack
	g.colour = colourRed
	if n == n.parent.left {
		g.rotateRight(t)
	} else {
		g.rotateLeft(t)
	}
}

func (n *node) rotateLeft(t *RedBlackTree) {
	right := n.right
	n.right = right.left
	if right.left != nil {
		right.left.parent = n
	}
	right.parent = n.parent
	switch {
	case n.parent == nil:
		t.root = right
	case n == n.parent.left:
		n.parent.left = right
	default:
		n.parent.right = right
	}
	right.left = n
	n.parent = right
}

func (n *node) rotateRight(t *RedBlackTree) {
	left := n.left
	n.left = left.right
	if left.right != nil {
		left.right.parent = n
	}
	left.parent = n.parent
	switch {
	case n.parent == nil:
		t.root = left
	case n == n.parent.right:
		n.parent.right = left
	default:
		n.parent.left = left
	}
	left.right = n
	n.parent = left
}

func (n *node) grandparent() *node {
	if n == nil || n.parent == nil {
		return nil
	}
	return n.parent.parent
}
<commit_msg>Update Insert -> Upsert<commit_after>// MIT License
//
// Copyright (c) 2017 Ryan Fowler
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

// Package tree provides an implementation of a red-black tree.
package tree

// Int represents an integer that implements the Item interface.
type Int int

// Less returns true if the Int is less than the provided Int. If the provided
// Item is not an Int, Less will panic.
func (i Int) Less(than Item) bool {
	return i < than.(Int)
}

// Item is the interface that wraps the Less method.
//
// Less should return 'true' if the instance is "less than" the provided Item.
// Items are considered equal if neither are less than each other.
// E.g. Items 'a' & 'b' are considered equal if: (!a.Less(b) && !b.Less(a))
type Item interface {
	Less(Item) bool
}

// RedBlackTree is an in-memory implementation of a red-black tree.
//
// The internal data structure will automatically re-balance, and therefore
// allow for O(log(n)) retrieval, insertion, and deletion.
//
// Note: While read-only operations may occur concurrently, any write operation
// must be serially executed (typically protected with a mutex).
type RedBlackTree struct {
	root *node
	size int
}

// Ascend (O(n)) starts at the first Item and calls 'fn' for each Item until no
// Items remain or fn returns 'false'.
func (t *RedBlackTree) Ascend(fn func(Item) bool) {
	if t.root == nil {
		return
	}
	n := t.root.min()
	for n != nil && fn(n.item) {
		n = n.next()
	}
}

// Descend (O(n)) starts at the last Item and calls 'fn' for each Item until no
// Items remain or fn returns 'false'.
func (t *RedBlackTree) Descend(fn func(Item) bool) {
	if t.root == nil {
		return
	}
	n := t.root.max()
	for n != nil && fn(n.item) {
		n = n.prev()
	}
}

// Delete (O(log(n))) deletes an item in the RedBlackTree equal to the provided
// item. If an item was deleted, it is returned. Otherwise, nil is returned.
//
// Note: equality for items a & b is: (!a.Less(b) && !b.Less(a)).
func (t *RedBlackTree) Delete(item Item) Item {
	if t.root == nil {
		return nil
	}
	return t.root.deleteItem(t, item)
}

// DeleteMax (O(log(n))) deletes the maximum item in the RedBlackTree, returning
// it. If the tree is empty, nil is returned.
func (t *RedBlackTree) DeleteMax() Item {
	if t.root == nil {
		return nil
	}
	return t.root.deleteMax(t)
}

// DeleteMin (O(log(n))) deletes the minimum item in the RedBlackTree, returning
// it. If the tree is empty, nil is returned.
func (t *RedBlackTree) DeleteMin() Item {
	if t.root == nil {
		return nil
	}
	return t.root.deleteMin(t)
}

// Get (O(log(n))) retrieves an item in the RedBlackTree equal to the provided
// item. If an item was found, it is returned. Otherwise, nil is returned.
//
// Note: equality for items a & b is: (!a.Less(b) && !b.Less(a)).
func (t *RedBlackTree) Get(item Item) Item {
	n := t.root.find(item)
	if n == nil {
		return nil
	}
	return n.item
}

// Upsert (O(log(n))) inserts (or replaces) an item into the RedBlackTree. If an
// item was replaced, it is returned. Otherwise, nil is returned.
//
// Note: equality for items a & b is: (!a.Less(b) && !b.Less(a)).
func (t *RedBlackTree) Upsert(item Item) Item {
	if t.root == nil {
		t.root = newNode(nil, item)
		t.root.colour = colourBlack
		t.size++
		return nil
	}
	n, oldItem := t.root.insert(item)
	if oldItem == nil {
		t.size++
		n.rebalanceInsert(t)
	}
	return oldItem
}

// Exists (O(log(n))) returns 'true' if an item equal to the provided item
// exists in the RedBlackTree.
//
// Note: equality for items a & b is: (!a.Less(b) && !b.Less(a)).
func (t *RedBlackTree) Exists(item Item) bool {
	return t.Get(item) != nil
}

// Min (O(log(n))) returns the minimum item in the RedBlackTree. If the tree is
// empty, nil is returned.
func (t *RedBlackTree) Min() Item {
	if t.root == nil {
		return nil
	}
	n := t.root
	for n.left != nil {
		n = n.left
	}
	return n.item
}

// Max (O(log(n))) returns the maximum item in the RedBlackTree. If the tree is
// empty, nil is returned.
func (t *RedBlackTree) Max() Item {
	if t.root == nil {
		return nil
	}
	n := t.root
	for n.right != nil {
		n = n.right
	}
	return n.item
}

// Size (O(1)) returns the number of items in the RedBlackTree.
func (t *RedBlackTree) Size() int {
	return t.size
}

type colour uint8

const (
	colourRed   colour = 0
	colourBlack colour = 1
)

type node struct {
	colour      colour
	parent      *node
	left, right *node
	item        Item
}

func newNode(parent *node, item Item) *node {
	return &node{
		colour: colourRed,
		parent: parent,
		item:   item,
	}
}

func (n *node) find(item Item) *node {
	for n != nil {
		switch {
		case item.Less(n.item):
			n = n.left
		case n.item.Less(item):
			n = n.right
		default:
			return n
		}
	}
	return nil
}

func (n *node) deleteMax(t *RedBlackTree) Item {
	return n.max().deleteNode(t)
}

func (n *node) deleteMin(t *RedBlackTree) Item {
	return n.min().deleteNode(t)
}

func (n *node) deleteItem(t *RedBlackTree, item Item) Item {
	n = n.find(item)
	if n == nil {
		return nil
	}
	return n.deleteNode(t)
}

func (n *node) deleteNode(t *RedBlackTree) Item {
	t.size--
	delItem := n.item

	var child, parent *node
	for {
		if n.left == nil {
			child = n.right
			parent = n.parent
			n.replaceNode(t, n.right)
			break
		}
		if n.right == nil {
			child = n.left
			parent = n.parent
			n.replaceNode(t, n.left)
			break
		}
		// replace minimum value in right subtree with node to delete.
		min := n.right.min()
		n.item = min.item
		n = min
	}

	if n.isRed() {
		return delItem
	}
	if child.isRed() {
		child.colour = colourBlack
		return delItem
	}
	child.rebalanceDelete(t, parent)
	return delItem
}

func (n *node) rebalanceDelete(t *RedBlackTree, parent *node) {
	var s *node
	for {
		// Case 1.
		if n == t.root {
			return
		}
		if n != nil {
			parent = n.parent
		}
		// Case 2.
		s = n.sibling(parent)
		if s.isRed() {
			parent.colour = colourRed
			s.colour = colourBlack
			if n == parent.left {
				parent.rotateLeft(t)
			} else {
				parent.rotateRight(t)
			}
		}
		// Case 3.
		s = n.sibling(parent)
		if parent.isBlack() && s.isBlack() && s != nil && s.left.isBlack() && s.right.isBlack() {
			s.colour = colourRed
			n = parent
			if n != nil {
				parent = n.parent
			} else {
				parent = nil
			}
			continue
		}
		break
	}
	// Case 4.
	if parent.isRed() &&
		s.isBlack() &&
		s != nil &&
		s.left.isBlack() &&
		s.right.isBlack() {
		s.colour = colourRed
		parent.colour = colourBlack
		return
	}
	// Case 5.
	if s.isBlack() && s != nil {
		if n == parent.left && s.right.isBlack() && s.left.isRed() {
			s.colour = colourRed
			s.left.colour = colourBlack
			s.rotateRight(t)
		} else if n == parent.right && s.left.isBlack() && s.right.isRed() {
			s.colour = colourRed
			s.right.colour = colourBlack
			s.rotateLeft(t)
		}
	}
	// Case 6.
	s = n.sibling(parent)
	if s != nil {
		s.colour = parent.colour
		parent.colour = colourBlack
		if n == parent.left {
			s.right.colour = colourBlack
			parent.rotateLeft(t)
		} else {
			s.left.colour = colourBlack
			parent.rotateRight(t)
		}
	}
}

func (n *node) isRed() bool {
	return n != nil && n.colour == colourRed
}

func (n *node) isBlack() bool {
	return n == nil || n.colour == colourBlack
}

func (n *node) sibling(parent *node) *node {
	if n == parent.left {
		return parent.right
	}
	return parent.left
}

func (n *node) replaceNode(t *RedBlackTree, child *node) {
	switch {
	case n.parent == nil:
		t.root = child
	case n == n.parent.left:
		n.parent.left = child
	default:
		n.parent.right = child
	}
	if child != nil {
		child.parent = n.parent
	}
}

func (n *node) min() *node {
	for n.left != nil {
		n = n.left
	}
	return n
}

func (n *node) max() *node {
	for n.right != nil {
		n = n.right
	}
	return n
}

func (n *node) next() *node {
	if n.right != nil {
		return n.right.min()
	}
	parent := n.parent
	for parent != nil && parent.right == n {
		n = parent
		parent = n.parent
	}
	return parent
}

func (n *node) prev() *node {
	if n.left != nil {
		return n.left.max()
	}
	parent := n.parent
	for parent != nil && parent.left == n {
		n = parent
		parent = n.parent
	}
	return parent
}

func (n *node) insert(item Item) (*node, Item) {
	for {
		switch {
		case item.Less(n.item):
			if n.left == nil {
				n.left = newNode(n, item)
				return n.left, nil
			}
			n = n.left
		case n.item.Less(item):
			if n.right == nil {
				n.right = newNode(n, item)
				return n.right, nil
			}
			n = n.right
		default:
			oldItem := n.item
			n.item = item
			return n, oldItem
		}
	}
}

func (n *node) rebalanceInsert(t *RedBlackTree) {
	var g *node
	for {
		// Case 1.
		if n.parent == nil {
			n.colour = colourBlack
			return
		}
		// Case 2.
		if n.parent.colour == colourBlack {
			return
		}
		// Case 3.
		g = n.grandparent()
		var ps *node
		if g != nil {
			if n.parent == g.left {
				ps = g.right
			} else {
				ps = g.left
			}
		}
		if ps.isBlack() {
			break
		}
		n.parent.colour = colourBlack
		ps.colour = colourBlack
		g.colour = colourRed
		n = g
	}
	// Case 4.
	if n == n.parent.right && n.parent == g.left {
		n.parent.rotateLeft(t)
		n = n.left
		g = n.grandparent()
	} else if n == n.parent.left && n.parent == g.right {
		n.parent.rotateRight(t)
		n = n.right
		g = n.grandparent()
	}
	// Case 5.
	n.parent.colour = colourBlack
	g.colour = colourRed
	if n == n.parent.left {
		g.rotateRight(t)
	} else {
		g.rotateLeft(t)
	}
}

func (n *node) rotateLeft(t *RedBlackTree) {
	right := n.right
	n.right = right.left
	if right.left != nil {
		right.left.parent = n
	}
	right.parent = n.parent
	switch {
	case n.parent == nil:
		t.root = right
	case n == n.parent.left:
		n.parent.left = right
	default:
		n.parent.right = right
	}
	right.left = n
	n.parent = right
}

func (n *node) rotateRight(t *RedBlackTree) {
	left := n.left
	n.left = left.right
	if left.right != nil {
		left.right.parent = n
	}
	left.parent = n.parent
	switch {
	case n.parent == nil:
		t.root = left
	case n == n.parent.right:
		n.parent.right = left
	default:
		n.parent.left = left
	}
	left.right = n
	n.parent = left
}

func (n *node) grandparent() *node {
	if n == nil || n.parent == nil {
		return nil
	}
	return n.parent.parent
}
<|endoftext|>


<commit_before>package model

// TODO: Use github.com/docker/libcompose

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"

	"github.com/docker/libcompose/config"
	"github.com/docker/libcompose/lookup"
	"github.com/docker/libcompose/project"
	"github.com/dtan4/paus-gitreceive/receiver/util"
	"github.com/pkg/errors"
	"gopkg.in/yaml.v2"
)

const (
	portBindingRegexp = `"?\d+:(\d+)"?`
)

var (
	portBinding = regexp.MustCompile(portBindingRegexp)
)

type Compose struct {
	ComposeFilePath string
	ProjectName     string

	dockerHost string
	project    *project.Project
}

type ComposeConfig struct {
	Version  string                           `yaml:"version,omitempty"`
	Services map[string]*config.ServiceConfig `yaml:"services,omitempty"`
	Volumes  map[string]*config.VolumeConfig  `yaml:"volumes,omitempty"`
	Networks map[string]*config.NetworkConfig `yaml:"networks,omitempty"`
}

func NewCompose(dockerHost, composeFilePath, projectName string) (*Compose, error) {
	ctx := project.Context{
		ComposeFiles: []string{composeFilePath},
		ProjectName:  projectName,
	}

	ctx.ResourceLookup = &lookup.FileConfigLookup{}
	ctx.EnvironmentLookup = &lookup.ComposableEnvLookup{
		Lookups: []config.EnvironmentLookup{
			&lookup.OsEnvLookup{},
		},
	}

	prj := project.NewProject(&ctx, nil, nil)

	if err := prj.Parse(); err != nil {
		return nil, errors.Wrap(err, "Failed to parse docker-compose.yml.")
	}

	return &Compose{
		ComposeFilePath: composeFilePath,
		ProjectName:     projectName,
		dockerHost:      dockerHost,
		project:         prj,
	}, nil
}

func (c *Compose) Build() error {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "build")
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)

	if err := util.RunCommand(cmd); err != nil {
		return err
	}

	return nil
}

func (c *Compose) GetContainerID(service string) (string, error) {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "ps", "-q", service)
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)
	out, err := cmd.Output()

	if err != nil {
		return "", errors.Wrapf(err, "Failed to get container ID. projectName: %s, service: %s", c.ProjectName, service)
	}

	return strings.Replace(string(out), "\n", "", -1), nil
}

func (c *Compose) InjectBuildArgs(buildArgs map[string]string) {
	webService := c.webService()

	if webService == nil {
		return
	}

	if webService.Build.Args == nil {
		webService.Build.Args = map[string]string{}
	}

	for k, v := range buildArgs {
		webService.Build.Args[k] = v
	}
}

func (c *Compose) InjectEnvironmentVariables(envs map[string]string) {
	webService := c.webService()

	if webService == nil {
		return
	}

	for k, v := range envs {
		webService.Environment = append(webService.Environment, fmt.Sprintf("%s=%s", k, v))
	}
}

func (c *Compose) Pull() error {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "pull")
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)

	if err := util.RunCommand(cmd); err != nil {
		return err
	}

	return nil
}

func (c *Compose) RewritePortBindings() {
	var newPorts []string

	for _, key := range c.project.ServiceConfigs.Keys() {
		if svc, ok := c.project.ServiceConfigs.Get(key); ok {
			if len(svc.Ports) == 0 {
				continue
			}

			newPorts = []string{}

			for _, port := range svc.Ports {
				matchResult := portBinding.FindStringSubmatch(port)

				if len(matchResult) == 2 {
					newPorts = append(newPorts, matchResult[1])
				} else {
					newPorts = append(newPorts, port)
				}
			}

			svc.Ports = newPorts
		}
	}
}

func (c *Compose) SaveAs(filePath string) error {
	services := map[string]*config.ServiceConfig{}

	for _, key := range c.project.ServiceConfigs.Keys() {
		if svc, ok := c.project.ServiceConfigs.Get(key); ok {
			services[key] = svc
		}
	}

	cfg := &ComposeConfig{
		Version:  "2",
		Services: services,
		Volumes:  c.project.VolumeConfigs,
		Networks: c.project.NetworkConfigs,
	}

	data, err := yaml.Marshal(cfg)

	if err != nil {
		return errors.Wrap(err, "Failed to generate YAML file.")
	}

	if err = ioutil.WriteFile(filePath, data, 0644); err != nil {
		return errors.Wrapf(err, "Failed to save as YAML file. path: %s", filePath)
	}

	c.ComposeFilePath = filePath

	return nil
}

func (c *Compose) Stop() error {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "stop")
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)

	if err := util.RunCommand(cmd); err != nil {
		return err
	}

	return nil
}

func (c *Compose) Up() error {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "up", "-d")
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)

	if err := util.RunCommand(cmd); err != nil {
		return err
	}

	return nil
}

func (c *Compose) webService() *config.ServiceConfig {
	if svc, ok := c.project.ServiceConfigs.Get("web"); ok {
		return svc
	}

	return nil
}
<commit_msg>Override existing environment variable<commit_after>package model

// TODO: Use github.com/docker/libcompose

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"

	"github.com/docker/libcompose/config"
	"github.com/docker/libcompose/lookup"
	"github.com/docker/libcompose/project"
	"github.com/dtan4/paus-gitreceive/receiver/util"
	"github.com/pkg/errors"
	"gopkg.in/yaml.v2"
)

const (
	portBindingRegexp = `"?\d+:(\d+)"?`
)

var (
	portBinding = regexp.MustCompile(portBindingRegexp)
)

type Compose struct {
	ComposeFilePath string
	ProjectName     string

	dockerHost string
	project    *project.Project
}

type ComposeConfig struct {
	Version  string                           `yaml:"version,omitempty"`
	Services map[string]*config.ServiceConfig `yaml:"services,omitempty"`
	Volumes  map[string]*config.VolumeConfig  `yaml:"volumes,omitempty"`
	Networks map[string]*config.NetworkConfig `yaml:"networks,omitempty"`
}

func NewCompose(dockerHost, composeFilePath, projectName string) (*Compose, error) {
	ctx := project.Context{
		ComposeFiles: []string{composeFilePath},
		ProjectName:  projectName,
	}

	ctx.ResourceLookup = &lookup.FileConfigLookup{}
	ctx.EnvironmentLookup = &lookup.ComposableEnvLookup{
		Lookups: []config.EnvironmentLookup{
			&lookup.OsEnvLookup{},
		},
	}

	prj := project.NewProject(&ctx, nil, nil)

	if err := prj.Parse(); err != nil {
		return nil, errors.Wrap(err, "Failed to parse docker-compose.yml.")
	}

	return &Compose{
		ComposeFilePath: composeFilePath,
		ProjectName:     projectName,
		dockerHost:      dockerHost,
		project:         prj,
	}, nil
}

func (c *Compose) Build() error {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "build")
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)

	if err := util.RunCommand(cmd); err != nil {
		return err
	}

	return nil
}

func (c *Compose) GetContainerID(service string) (string, error) {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "ps", "-q", service)
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)
	out, err := cmd.Output()

	if err != nil {
		return "", errors.Wrapf(err, "Failed to get container ID. projectName: %s, service: %s", c.ProjectName, service)
	}

	return strings.Replace(string(out), "\n", "", -1), nil
}

func (c *Compose) InjectBuildArgs(buildArgs map[string]string) {
	webService := c.webService()

	if webService == nil {
		return
	}

	if webService.Build.Args == nil {
		webService.Build.Args = map[string]string{}
	}

	for k, v := range buildArgs {
		webService.Build.Args[k] = v
	}
}

func (c *Compose) InjectEnvironmentVariables(envs map[string]string) {
	webService := c.webService()

	if webService == nil {
		return
	}

	envmap := make(map[string]string)

	for _, env := range webService.Environment {
		kv := strings.SplitN(env, "=", 2)
		envmap[kv[0]] = kv[1]
	}

	for k, v := range envs {
		envmap[k] = v
	}

	webService.Environment = []string{}

	for k, v := range envmap {
		webService.Environment = append(webService.Environment, fmt.Sprintf("%s=%s", k, v))
	}
}

func (c *Compose) Pull() error {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "pull")
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)

	if err := util.RunCommand(cmd); err != nil {
		return err
	}

	return nil
}

func (c *Compose) RewritePortBindings() {
	var newPorts []string

	for _, key := range c.project.ServiceConfigs.Keys() {
		if svc, ok := c.project.ServiceConfigs.Get(key); ok {
			if len(svc.Ports) == 0 {
				continue
			}

			newPorts = []string{}

			for _, port := range svc.Ports {
				matchResult := portBinding.FindStringSubmatch(port)

				if len(matchResult) == 2 {
					newPorts = append(newPorts, matchResult[1])
				} else {
					newPorts = append(newPorts, port)
				}
			}

			svc.Ports = newPorts
		}
	}
}

func (c *Compose) SaveAs(filePath string) error {
	services := map[string]*config.ServiceConfig{}

	for _, key := range c.project.ServiceConfigs.Keys() {
		if svc, ok := c.project.ServiceConfigs.Get(key); ok {
			services[key] = svc
		}
	}

	cfg := &ComposeConfig{
		Version:  "2",
		Services: services,
		Volumes:  c.project.VolumeConfigs,
		Networks: c.project.NetworkConfigs,
	}

	data, err := yaml.Marshal(cfg)

	if err != nil {
		return errors.Wrap(err, "Failed to generate YAML file.")
	}

	if err = ioutil.WriteFile(filePath, data, 0644); err != nil {
		return errors.Wrapf(err, "Failed to save as YAML file. path: %s", filePath)
	}

	c.ComposeFilePath = filePath

	return nil
}

func (c *Compose) Stop() error {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "stop")
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)

	if err := util.RunCommand(cmd); err != nil {
		return err
	}

	return nil
}

func (c *Compose) Up() error {
	cmd := exec.Command("docker-compose", "-f", c.ComposeFilePath, "-p", c.ProjectName, "up", "-d")
	cmd.Env = append(os.Environ(), "DOCKER_HOST="+c.dockerHost)

	if err := util.RunCommand(cmd); err != nil {
		return err
	}

	return nil
}

func (c *Compose) webService() *config.ServiceConfig {
	if svc, ok := c.project.ServiceConfigs.Get("web"); ok {
		return svc
	}

	return nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"context"
	"crypto/tls"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/cybozu-go/log"
)

const (
	defaultHTTPReadTimeout = 30 * time.Second

	// our request tracking header.
	uuidHeaderName = "X-Cybozu-Request-ID"
)

// HTTPServer is a wrapper for http.Server.
//
// This struct overrides Serve and ListenAndServe* methods.
//
// http.Server members are replaced as following:
//    - Handler is replaced with a wrapper handler.
//    - ReadTimeout is set to 30 seconds if it is zero.
//    - ConnState is replaced with the one provided by the framework.
type HTTPServer struct {
	*http.Server

	// AccessLog is a logger for access logs.
	// If this is nil, the default logger is used.
	AccessLog *log.Logger

	// ShutdownTimeout is the maximum duration the server waits for
	// all connections to be closed before shutdown.
	//
	// Zero duration disables timeout.
	ShutdownTimeout time.Duration

	// Env is the environment where this server runs.
	//
	// The global environment is used if Env is nil.
	Env *Environment

	handler  http.Handler
	wg       sync.WaitGroup
	initOnce sync.Once
}

type logResponseWriter struct {
	http.ResponseWriter
	status int
	size   int
}

func (w *logResponseWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *logResponseWriter) Write(data []byte) (int, error) {
	n, err := w.ResponseWriter.Write(data)
	w.size += n
	return n, err
}

// ServeHTTP implements http.Handler interface.
func (s *HTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	startTime := time.Now()

	lw := &logResponseWriter{w, http.StatusOK, 0}
	ctx, cancel := context.WithCancel(s.Env.ctx)
	defer cancel()
	s.handler.ServeHTTP(lw, r.WithContext(ctx))

	fields := map[string]interface{}{
		log.FnAccessLog:      true,
		log.FnResponseTime:   time.Since(startTime).Seconds(),
		log.FnProtocol:       r.Proto,
		log.FnHTTPStatusCode: lw.status,
		log.FnHTTPMethod:     r.Method,
		log.FnURL:            r.RequestURI,
		log.FnHTTPHost:       r.Host,
		log.FnRequestSize:    r.ContentLength,
		log.FnResponseSize:   lw.size,
	}
	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		fields[log.FnRemoteAddress] = ip
	}
	ua := r.Header.Get("User-Agent")
	if len(ua) > 0 {
		fields[log.FnHTTPUserAgent] = ua
	}
	reqid := r.Header.Get(uuidHeaderName)
	if len(reqid) > 0 {
		fields[log.FnRequestID] = reqid
	}

	lv := log.LvInfo
	switch {
	case 500 <= lw.status:
		lv = log.LvError
	case 400 <= lw.status:
		lv = log.LvWarn
	}
	s.AccessLog.Log(lv, "cmd: "+http.StatusText(lw.status), fields)
}

func (s *HTTPServer) init() {
	if s.handler != nil {
		return
	}

	if s.Server.Handler == nil {
		panic("Handler must not be nil")
	}
	s.handler = s.Server.Handler
	s.Server.Handler = s
	if s.Server.ReadTimeout == 0 {
		s.Server.ReadTimeout = defaultHTTPReadTimeout
	}
	s.Server.ConnState = func(c net.Conn, state http.ConnState) {
		if state == http.StateNew {
			s.wg.Add(1)
			return
		}
		if state == http.StateHijacked || state == http.StateClosed {
			s.wg.Done()
		}
	}

	if s.AccessLog == nil {
		s.AccessLog = log.DefaultLogger()
	}

	if s.Env == nil {
		s.Env = defaultEnv
	}
	s.Env.Go(s.wait)
}

func (s *HTTPServer) wait(ctx context.Context) error {
	<-ctx.Done()

	// shorten keep-alive timeout
	s.Server.ReadTimeout = 100 * time.Millisecond
	s.Server.SetKeepAlivesEnabled(false)

	if s.ShutdownTimeout == 0 {
		s.wg.Wait()
		return nil
	}

	ch := make(chan struct{})
	go func() {
		s.wg.Wait()
		close(ch)
	}()

	select {
	case <-ch:
	case <-time.After(s.ShutdownTimeout):
		log.Warn("cmd: timeout waiting for shutdown", nil)
	}
	return nil
}

// Serve overrides http.Server's Serve method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.
//
// The framework automatically closes l when the environment's Stop
// is called.
//
// Serve always returns nil.
func (s *HTTPServer) Serve(l net.Listener) error {
	s.initOnce.Do(s.init)

	go func() {
		<-s.Env.ctx.Done()
		l.Close()
	}()

	go func() {
		s.Server.Serve(l)
	}()

	return nil
}

type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}

// ListenAndServe overrides http.Server's method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.
//
// ListenAndServe returns non-nil error if and only if net.Listen failed.
func (s *HTTPServer) ListenAndServe() error {
	addr := s.Server.Addr
	if addr == "" {
		addr = ":http"
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return s.Serve(tcpKeepAliveListener{ln.(*net.TCPListener)})
}

// ListenAndServeTLS overrides http.Server's method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.
//
// Another difference from the original is that certFile and keyFile
// must be specified.  If not, configure http.Server.TLSConfig
// manually and use Serve().
//
// HTTP/2 is always enabled.
//
// ListenAndServeTLS returns non-nil error if net.Listen failed
// or failed to load certificate files.
func (s *HTTPServer) ListenAndServeTLS(certFile, keyFile string) error {
	addr := s.Server.Addr
	if addr == "" {
		addr = ":https"
	}

	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	config := &tls.Config{
		NextProtos:               []string{"h2", "http/1.1"},
		Certificates:             []tls.Certificate{cert},
		PreferServerCipherSuites: true,
		ClientSessionCache:       tls.NewLRUClientSessionCache(0),
	}
	s.Server.TLSConfig = config

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, config)
	return s.Serve(tlsListener)
}
<commit_msg>misc.<commit_after>package cmd

import (
	"context"
	"crypto/tls"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/cybozu-go/log"
)

const (
	defaultHTTPReadTimeout = 30 * time.Second

	// our request tracking header.
	uuidHeaderName = "X-Cybozu-Request-ID"
)

// HTTPServer is a wrapper for http.Server.
//
// This struct overrides Serve and ListenAndServe* methods.
//
// http.Server members are replaced as following:
//    - Handler is replaced with a wrapper handler.
//    - ReadTimeout is set to 30 seconds if it is zero.
//    - ConnState is replaced with the one provided by the framework.
type HTTPServer struct {
	*http.Server

	// AccessLog is a logger for access logs.
	// If this is nil, the default logger is used.
	AccessLog *log.Logger

	// ShutdownTimeout is the maximum duration the server waits for
	// all connections to be closed before shutdown.
	//
	// Zero duration disables timeout.
	ShutdownTimeout time.Duration

	// Env is the environment where this server runs.
	//
	// The global environment is used if Env is nil.
	Env *Environment

	handler  http.Handler
	wg       sync.WaitGroup
	initOnce sync.Once
}

type logResponseWriter struct {
	http.ResponseWriter
	status int
	size   int
}

func (w *logResponseWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *logResponseWriter) Write(data []byte) (int, error) {
	n, err := w.ResponseWriter.Write(data)
	w.size += n
	return n, err
}

// ServeHTTP implements http.Handler interface.
func (s *HTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	startTime := time.Now()

	lw := &logResponseWriter{w, http.StatusOK, 0}
	ctx, cancel := context.WithCancel(s.Env.ctx)
	defer cancel()
	s.handler.ServeHTTP(lw, r.WithContext(ctx))

	fields := map[string]interface{}{
		log.FnType:           "access",
		log.FnResponseTime:   time.Since(startTime).Seconds(),
		log.FnProtocol:       r.Proto,
		log.FnHTTPStatusCode: lw.status,
		log.FnHTTPMethod:     r.Method,
		log.FnURL:            r.RequestURI,
		log.FnHTTPHost:       r.Host,
		log.FnRequestSize:    r.ContentLength,
		log.FnResponseSize:   lw.size,
	}
	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		fields[log.FnRemoteAddress] = ip
	}
	ua := r.Header.Get("User-Agent")
	if len(ua) > 0 {
		fields[log.FnHTTPUserAgent] = ua
	}
	reqid := r.Header.Get(uuidHeaderName)
	if len(reqid) > 0 {
		fields[log.FnRequestID] = reqid
	}

	lv := log.LvInfo
	switch {
	case 500 <= lw.status:
		lv = log.LvError
	case 400 <= lw.status:
		lv = log.LvWarn
	}
	s.AccessLog.Log(lv, "cmd: "+http.StatusText(lw.status), fields)
}

func (s *HTTPServer) init() {
	if s.handler != nil {
		return
	}

	if s.Server.Handler == nil {
		panic("Handler must not be nil")
	}
	s.handler = s.Server.Handler
	s.Server.Handler = s
	if s.Server.ReadTimeout == 0 {
		s.Server.ReadTimeout = defaultHTTPReadTimeout
	}
	s.Server.ConnState = func(c net.Conn, state http.ConnState) {
		if state == http.StateNew {
			s.wg.Add(1)
			return
		}
		if state == http.StateHijacked || state == http.StateClosed {
			s.wg.Done()
		}
	}

	if s.AccessLog == nil {
		s.AccessLog = log.DefaultLogger()
	}

	if s.Env == nil {
		s.Env = defaultEnv
	}
	s.Env.Go(s.wait)
}

func (s *HTTPServer) wait(ctx context.Context) error {
	<-ctx.Done()

	// shorten keep-alive timeout
	s.Server.ReadTimeout = 100 * time.Millisecond
	s.Server.SetKeepAlivesEnabled(false)

	if s.ShutdownTimeout == 0 {
		s.wg.Wait()
		return nil
	}

	ch := make(chan struct{})
	go func() {
		s.wg.Wait()
		close(ch)
	}()

	select {
	case <-ch:
	case <-time.After(s.ShutdownTimeout):
		log.Warn("cmd: timeout waiting for shutdown", nil)
	}
	return nil
}

// Serve overrides http.Server's Serve method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.
//
// The framework automatically closes l when the environment's Stop
// is called.
//
// Serve always returns nil.
func (s *HTTPServer) Serve(l net.Listener) error {
	s.initOnce.Do(s.init)

	go func() {
		<-s.Env.ctx.Done()
		l.Close()
	}()

	go func() {
		s.Server.Serve(l)
	}()

	return nil
}

type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}

// ListenAndServe overrides http.Server's method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.
//
// ListenAndServe returns non-nil error if and only if net.Listen failed.
func (s *HTTPServer) ListenAndServe() error {
	addr := s.Server.Addr
	if addr == "" {
		addr = ":http"
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return s.Serve(tcpKeepAliveListener{ln.(*net.TCPListener)})
}

// ListenAndServeTLS overrides http.Server's method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.
//
// Another difference from the original is that certFile and keyFile
// must be specified.  If not, configure http.Server.TLSConfig
// manually and use Serve().
//
// HTTP/2 is always enabled.
//
// ListenAndServeTLS returns non-nil error if net.Listen failed
// or failed to load certificate files.
func (s *HTTPServer) ListenAndServeTLS(certFile, keyFile string) error {
	addr := s.Server.Addr
	if addr == "" {
		addr = ":https"
	}

	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	config := &tls.Config{
		NextProtos:               []string{"h2", "http/1.1"},
		Certificates:             []tls.Certificate{cert},
		PreferServerCipherSuites: true,
		ClientSessionCache:       tls.NewLRUClientSessionCache(0),
	}
	s.Server.TLSConfig = config

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, config)
	return s.Serve(tlsListener)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
)

func (fw *Flywheel) SendPing(start bool) Pong {
	replyTo := make(chan Pong, 1)
	sreq := Ping{replyTo: replyTo, requestStart: start}

	fw.pings <- sreq

	status := <-replyTo
	return status
}

func (fw *Flywheel) ProxyEndpoint(hostname string) string {
	vhost, ok := fw.config.Vhosts[hostname]
	if ok {
		return vhost
	}
	return fw.config.Endpoint
}

func (fw *Flywheel) Proxy(w http.ResponseWriter, r *http.Request) {
	client := &http.Client{}
	r.URL.Query().Del("flywheel")

	endpoint := fw.ProxyEndpoint(r.Host)
	if endpoint == "" {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("Invalid flywheel endpoint config"))
		log.Fatal("Invalid endpoint URL")
	}

	r.URL.Scheme = "http"

	r.URL.Host = endpoint
	r.RequestURI = ""
	resp, err := client.Do(r)
	if err != nil {
		log.Print(err)
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}

	for key, value := range resp.Header {
		w.Header()[key] = value
	}
	w.WriteHeader(resp.StatusCode)

	_, err = io.Copy(w, resp.Body)
	if err != nil {
		log.Print(err)
	}
}

func (fw *Flywheel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Printf("[%s] %s %s", r.RemoteAddr, r.Method, r.RequestURI)

	query := r.URL.Query()
	flywheel, ok := query["flywheel"]
	pong := fw.SendPing(ok && flywheel[0] == "start")

	if ok {
		query.Del("flywheel")
		r.URL.RawQuery = query.Encode()
		w.Header().Set("Location", r.URL.String())
		w.WriteHeader(302)
		return
	}

	if pong.Err != nil {
		body := fmt.Sprintf(HTML_ERROR, pong.Err)
		w.Write([]byte(body))
		return
	}

	switch pong.Status {
	case STOPPED:
		query.Set("flywheel", "start")
		r.URL.RawQuery = query.Encode()
		body := fmt.Sprintf(HTML_STOPPED, r.URL)
		w.Write([]byte(body))
	case STARTING:
		w.Write([]byte(HTML_STARTING))
	case STARTED:
		fw.Proxy(w, r)
	case STOPPING:
		w.Write([]byte(HTML_STOPPING))
	case UNHEALTHY:
		w.Write([]byte(HTML_UNHEALTHY))
	}
}
<commit_msg>Set non-200 status codes when powered down, etc<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
)

func (fw *Flywheel) SendPing(start bool) Pong {
	replyTo := make(chan Pong, 1)
	sreq := Ping{replyTo: replyTo, requestStart: start}

	fw.pings <- sreq

	status := <-replyTo
	return status
}

func (fw *Flywheel) ProxyEndpoint(hostname string) string {
	vhost, ok := fw.config.Vhosts[hostname]
	if ok {
		return vhost
	}
	return fw.config.Endpoint
}

func (fw *Flywheel) Proxy(w http.ResponseWriter, r *http.Request) {
	client := &http.Client{}
	r.URL.Query().Del("flywheel")

	endpoint := fw.ProxyEndpoint(r.Host)
	if endpoint == "" {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("Invalid flywheel endpoint config"))
		log.Fatal("Invalid endpoint URL")
	}

	r.URL.Scheme = "http"

	r.URL.Host = endpoint
	r.RequestURI = ""
	resp, err := client.Do(r)
	if err != nil {
		log.Print(err)
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}

	for key, value := range resp.Header {
		w.Header()[key] = value
	}
	w.WriteHeader(resp.StatusCode)

	_, err = io.Copy(w, resp.Body)
	if err != nil {
		log.Print(err)
	}
}

func (fw *Flywheel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Printf("[%s] %s %s", r.RemoteAddr, r.Method, r.RequestURI)

	query := r.URL.Query()
	flywheel, ok := query["flywheel"]
	pong := fw.SendPing(ok && flywheel[0] == "start")

	if ok {
		query.Del("flywheel")
		r.URL.RawQuery = query.Encode()
		w.Header().Set("Location", r.URL.String())
		w.WriteHeader(http.StatusTemporaryRedirect)
		return
	}

	if pong.Err != nil {
		body := fmt.Sprintf(HTML_ERROR, pong.Err)
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(body))
		return
	}

	switch pong.Status {
	case STOPPED:
		query.Set("flywheel", "start")
		r.URL.RawQuery = query.Encode()
		body := fmt.Sprintf(HTML_STOPPED, r.URL)
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(body))
	case STARTING:
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(HTML_STARTING))
	case STARTED:
		fw.Proxy(w, r)
	case STOPPING:
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(HTML_STOPPING))
	case UNHEALTHY:
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(HTML_UNHEALTHY))
	}
}
<|endoftext|>


<commit_before>package pilosa

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
)

const MAX_QUERIES = 1000

// Client is not safe for concurrent usage.
type Client struct {
	pilosaURL string
	queries   []string
}

func NewClient(pilosaURL string) *Client {
	return &Client{
		pilosaURL: pilosaURL,
		queries:   make([]string, 0),
	}
}

func (c *Client) AddQuery(query string) {
	c.queries = append(c.queries, query)
}

type Results struct {
	Results []interface{}
}

func (c *Client) ExecuteQueries(db int) (Results, error) {
	if len(c.queries) == 0 {
		return Results{}, nil
	}
	r := Results{}
	err := c.pilosaPost(bytes.NewBufferString(strings.Join(c.queries, "")), db, &r)
	return r, err
}

func (c *Client) ClearQueries() {
	c.queries = c.queries[:0]
}

type SetBitResponse struct {
	Results []bool
}

func (c *Client) SetBit(db int, bitmapID int, frame string, profileID int) (bool, error) {
	query := bytes.NewBufferString(fmt.Sprintf("SetBit(%d, '%s', %d)", bitmapID, frame, profileID))
	resp := SetBitResponse{}
	err := c.pilosaPost(query, db, &resp)
	if err != nil {
		return false, err
	}
	if len(resp.Results) != 1 {
		return false, fmt.Errorf("Unexpected response from SetBit: %v", resp)
	}
	return resp.Results[0], nil
}

type ClearBitResponse struct {
	Results []bool
}

func (c *Client) ClearBit(db int, bitmapID int, frame string, profileID int) (bool, error) {
	query := bytes.NewBufferString(fmt.Sprintf("ClearBit(%d, '%s', %d)", bitmapID, frame, profileID))
	resp := ClearBitResponse{}
	err := c.pilosaPost(query, db, &resp)
	if err != nil {
		return false, err
	}
	if len(resp.Results) != 1 {
		return false, fmt.Errorf("Unexpected response from ClearBit: %v", resp)
	}
	return resp.Results[0], nil
}

type CountBitResponse struct {
	Results []int64
}

func (c *Client) CountBit(db int, bitmapID int, frame string) (int64, error) {
	query := bytes.NewBufferString(fmt.Sprintf("Count(Bitmap(%d, '%s'))", bitmapID, frame))
	resp := CountBitResponse{}
	err := c.pilosaPost(query, db, &resp)
	if err != nil {
		return 0, err
	}
	if len(resp.Results) != 1 {
		return 0, fmt.Errorf("Unexpected response from CountBit: %v", resp)
	}
	return resp.Results[0], nil
}

func (c *Client) pilosaPostRaw(query io.Reader, db int) (string, error) {
	postURL := fmt.Sprintf("%s/query?db=%d", c.pilosaURL, db)
	req, err := http.Post(postURL, "application/pql", query)
	if err != nil {
		return "", err
	}

	buf, err := ioutil.ReadAll(req.Body)
	return string(buf), err
}

func (c *Client) pilosaPost(query io.Reader, db int, v interface{}) error {
	postURL := fmt.Sprintf("%s/query?db=%d", c.pilosaURL, db)
	req, err := http.Post(postURL, "application/pql", query)
	if err != nil {
		return fmt.Errorf("error with http.Post in pilosaPost: %v", err)
	}
	if req.StatusCode >= 400 {
		bod, err := ioutil.ReadAll(req.Body)
		if err != nil {
			bod = []byte("")
		}
		return fmt.Errorf("bad status: %v - Body: %v", req.Status, string(bod))
	}
	dec := json.NewDecoder(req.Body)

	err = dec.Decode(v)
	if err != nil {
		return fmt.Errorf("error json decoding request body: %v", err)
	}
	return nil

}
<commit_msg>db should be string not int<commit_after>package pilosa

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
)

const MAX_QUERIES = 1000

// Client is not safe for concurrent usage.
type Client struct {
	pilosaURL string
	queries   []string
}

func NewClient(pilosaURL string) *Client {
	return &Client{
		pilosaURL: pilosaURL,
		queries:   make([]string, 0),
	}
}

func (c *Client) AddQuery(query string) {
	c.queries = append(c.queries, query)
}

type Results struct {
	Results []interface{}
}

func (c *Client) ExecuteQueries(db string) (Results, error) {
	if len(c.queries) == 0 {
		return Results{}, nil
	}
	r := Results{}
	err := c.pilosaPost(bytes.NewBufferString(strings.Join(c.queries, "")), db, &r)
	return r, err
}

func (c *Client) ClearQueries() {
	c.queries = c.queries[:0]
}

type SetBitResponse struct {
	Results []bool
}

func (c *Client) SetBit(db string, bitmapID int, frame string, profileID int) (bool, error) {
	query := bytes.NewBufferString(fmt.Sprintf("SetBit(%d, '%s', %d)", bitmapID, frame, profileID))
	resp := SetBitResponse{}
	err := c.pilosaPost(query, db, &resp)
	if err != nil {
		return false, err
	}
	if len(resp.Results) != 1 {
		return false, fmt.Errorf("Unexpected response from SetBit: %v", resp)
	}
	return resp.Results[0], nil
}

type ClearBitResponse struct {
	Results []bool
}

func (c *Client) ClearBit(db string, bitmapID int, frame string, profileID int) (bool, error) {
	query := bytes.NewBufferString(fmt.Sprintf("ClearBit(%d, '%s', %d)", bitmapID, frame, profileID))
	resp := ClearBitResponse{}
	err := c.pilosaPost(query, db, &resp)
	if err != nil {
		return false, err
	}
	if len(resp.Results) != 1 {
		return false, fmt.Errorf("Unexpected response from ClearBit: %v", resp)
	}
	return resp.Results[0], nil
}

type CountBitResponse struct {
	Results []int64
}

func (c *Client) CountBit(db string, bitmapID int, frame string) (int64, error) {
	query := bytes.NewBufferString(fmt.Sprintf("Count(Bitmap(%d, '%s'))", bitmapID, frame))
	resp := CountBitResponse{}
	err := c.pilosaPost(query, db, &resp)
	if err != nil {
		return 0, err
	}
	if len(resp.Results) != 1 {
		return 0, fmt.Errorf("Unexpected response from CountBit: %v", resp)
	}
	return resp.Results[0], nil
}

func (c *Client) pilosaPostRaw(query io.Reader, db string) (string, error) {
	postURL := fmt.Sprintf("%s/query?db=%s", c.pilosaURL, db)
	req, err := http.Post(postURL, "application/pql", query)
	if err != nil {
		return "", err
	}

	buf, err := ioutil.ReadAll(req.Body)
	return string(buf), err
}

func (c *Client) pilosaPost(query io.Reader, db string, v interface{}) error {
	postURL := fmt.Sprintf("%s/query?db=%s", c.pilosaURL, db)
	req, err := http.Post(postURL, "application/pql", query)
	if err != nil {
		return fmt.Errorf("error with http.Post in pilosaPost: %v", err)
	}
	if req.StatusCode >= 400 {
		bod, err := ioutil.ReadAll(req.Body)
		if err != nil {
			bod = []byte("")
		}
		return fmt.Errorf("bad status: %v - Body: %v", req.Status, string(bod))
	}
	dec := json.NewDecoder(req.Body)

	err = dec.Decode(v)
	if err != nil {
		return fmt.Errorf("error json decoding request body: %v", err)
	}
	return nil

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math"

	"github.com/jzacsh/netwtcpip-cmp405/parseip4"
)

var partTwoHosts = []parseip4.Addr{
	{IP: parseip4.NewAddr(9, 201, 195, 84), Mask: parseip4.NewAddr(255, 255, 240, 0)},
	{IP: parseip4.NewAddr(128, 10, 189, 215), Mask: parseip4.NewAddr(255, 255, 248, 0)},
	{IP: parseip4.NewAddr(135, 21, 243, 82), Mask: parseip4.NewAddr(255, 255, 224, 0)},
	{IP: parseip4.NewAddr(75, 149, 205, 61), Mask: parseip4.NewAddr(255, 255, 192, 0)},
	{IP: parseip4.NewAddr(7, 105, 198, 111), Mask: parseip4.NewAddr(255, 255, 252, 0)},

	// TODO(zacsh) remove this sample from the last slide
	{IP: parseip4.NewAddr(128, 10, 211, 78), Mask: parseip4.NewAddr(255, 255, 240, 0)},
}

type subnetRequisites struct {
	ClassfulContext parseip4.OctsList
	MaxSubnets      uint
	SubnetIndex     parseip4.Octets
	HostIndex       parseip4.Octets
}

type OptimalSubnet struct {
	MinSubnetBits     uint
	MaxHostsPerSubnet parseip4.Octets
	Address           parseip4.Addr
}

var partOneGivens = []subnetRequisites{
	{parseip4.OctsList{128, 10, 0, 0}, 55, 51, 121},
	{parseip4.OctsList{128, 10, 0, 0}, 55, 42, 867},
	{parseip4.OctsList{128, 10, 0, 0}, 121, 115, 246},
	{parseip4.OctsList{128, 10, 0, 0}, 121, 97, 443},
	{parseip4.OctsList{128, 10, 0, 0}, 26, 19, 237},
	{parseip4.OctsList{128, 10, 0, 0}, 26, 25, 1397},
	{parseip4.OctsList{128, 10, 0, 0}, 261, 227, 86},
	{parseip4.OctsList{128, 10, 0, 0}, 261, 259, 49},
	{parseip4.OctsList{128, 10, 0, 0}, 529, 519, 33},
	{parseip4.OctsList{128, 10, 0, 0}, 529, 510, 59},
}

func (s *subnetRequisites) String() string {
	return fmt.Sprintf(
		"max subnets: %d, subnet index: %d, host index: %d",
		s.MaxSubnets, s.SubnetIndex, s.HostIndex)
}

func maxIntWithBits(nbits uint) uint32 {
	// 1 because 2^N bits only gets 2^N-1 if all 1s. +1 more because all 1s is
	// reserved for broadcast.
	const gap float64 = 2

	maxInt := math.Pow(2, float64(nbits))
	if maxInt < gap {
		// we want to avoid underflows, so stick to the point of the API and return
		// effectively zero
		return 0
	}

	return uint32(maxInt - gap)
}

func (s *subnetRequisites) FindSolution() OptimalSubnet {
	opt := OptimalSubnet{}
	_, classCidrOffset, _ := parseip4.Classful(s.ClassfulContext)

	// Brute force solve for Ceil(log2(s.MaxSubnets))
	for {
		if maxIntWithBits(opt.MinSubnetBits) >= uint32(s.MaxSubnets) {
			break
		}
		opt.MinSubnetBits++
	}

	opt.MaxHostsPerSubnet = parseip4.Octets(maxIntWithBits(32 - classCidrOffset - opt.MinSubnetBits))

	mask := parseip4.Octets(0xFFFFFFFF)
	mask <<= (32 - opt.MinSubnetBits)
	opt.Address.Mask = mask.List()

	subnetBitCount := parseip4.CountBitSize(s.SubnetIndex)
	hostBitAddrSpace := 32 - classCidrOffset - subnetBitCount

	ip := s.ClassfulContext.Pack() |
		parseip4.Octets(s.SubnetIndex<<hostBitAddrSpace) |
		s.HostIndex
	opt.Address.IP = ip.List()

	return opt
}

func main() {
	fmt.Printf("part 1: analyzing %d hosts ...\n", len(partOneGivens))
	for _, req := range partOneGivens {
		sol := req.FindSolution()
		fmt.Printf(
			"  given: %s\n\tmin # of subnet bits: %d\n\tmax # hosts per subnet: %d\n\taddress: %s\n",
			req.String(),
			sol.MinSubnetBits,
			sol.MaxHostsPerSubnet,
			sol.Address.String())
	}

	fmt.Printf("\npart 2: analyzing %d hosts ...\n", len(partTwoHosts))
	for _, addr := range partTwoHosts {
		classMask, _, klass := parseip4.Classful(addr.IP)

		fmt.Printf(
			"  network: %v (class %s masked)\n\t%v\n\tnetwork id:\t%d\n\t subnet id:\t%d\n\t   host id:\t%d\n",
			(addr.IP.Pack() & classMask.Pack()).List(), klass,
			addr.String(),
			addr.NetworkIndex(),
			addr.SubnetIndex(),
			addr.HostIndex())
	}
}
<commit_msg>noop(doc) added sanity-checks from slide deck; all perfect<commit_after>package main

import (
	"fmt"
	"math"

	"github.com/jzacsh/netwtcpip-cmp405/parseip4"
)

var partTwoHosts = []parseip4.Addr{
	{IP: parseip4.NewAddr(9, 201, 195, 84), Mask: parseip4.NewAddr(255, 255, 240, 0)},
	{IP: parseip4.NewAddr(128, 10, 189, 215), Mask: parseip4.NewAddr(255, 255, 248, 0)},
	{IP: parseip4.NewAddr(135, 21, 243, 82), Mask: parseip4.NewAddr(255, 255, 224, 0)},
	{IP: parseip4.NewAddr(75, 149, 205, 61), Mask: parseip4.NewAddr(255, 255, 192, 0)},
	{IP: parseip4.NewAddr(7, 105, 198, 111), Mask: parseip4.NewAddr(255, 255, 252, 0)},

	// TODO(zacsh) remove this sample from the last slide
	{IP: parseip4.NewAddr(128, 10, 211, 78), Mask: parseip4.NewAddr(255, 255, 240, 0)},
}

type subnetRequisites struct {
	ClassfulContext parseip4.OctsList
	MaxSubnets      uint
	SubnetIndex     parseip4.Octets
	HostIndex       parseip4.Octets
}

type OptimalSubnet struct {
	MinSubnetBits     uint
	MaxHostsPerSubnet parseip4.Octets
	Address           parseip4.Addr
}

var partOneGivens = []subnetRequisites{
	{parseip4.OctsList{128, 10, 0, 0}, 55, 51, 121},
	{parseip4.OctsList{128, 10, 0, 0}, 55, 42, 867},
	{parseip4.OctsList{128, 10, 0, 0}, 121, 115, 246},
	{parseip4.OctsList{128, 10, 0, 0}, 121, 97, 443},
	{parseip4.OctsList{128, 10, 0, 0}, 26, 19, 237},
	{parseip4.OctsList{128, 10, 0, 0}, 26, 25, 1397},
	{parseip4.OctsList{128, 10, 0, 0}, 261, 227, 86},
	{parseip4.OctsList{128, 10, 0, 0}, 261, 259, 49},
	{parseip4.OctsList{128, 10, 0, 0}, 529, 519, 33},
	{parseip4.OctsList{128, 10, 0, 0}, 529, 510, 59},
}

// Sanity checks taken from slides:
// http://comet.lehman.cuny.edu/sfakhouri/teaching/cmp/cmp405/f17/classnotes/Subnets%20To%20IP%20Address.pdf
var partOneSanityCheck = []subnetRequisites{
	{parseip4.OctsList{128, 10, 0, 0}, 200, 193, 129}, // slide 1
	{parseip4.OctsList{128, 10, 0, 0}, 120, 119, 68},  // slide 2
	{parseip4.OctsList{128, 10, 0, 0}, 120, 105, 352}, // slide 3
	{parseip4.OctsList{128, 10, 0, 0}, 58, 45, 98},    // slide 4
	{parseip4.OctsList{128, 10, 0, 0}, 58, 48, 598},   // slide 5
	{parseip4.OctsList{128, 10, 0, 0}, 29, 28, 59},    // slide 6
	{parseip4.OctsList{128, 10, 0, 0}, 29, 25, 1069},  // slide 7
}

func (s *subnetRequisites) String() string {
	return fmt.Sprintf(
		"max subnets: %d, subnet index: %d, host index: %d",
		s.MaxSubnets, s.SubnetIndex, s.HostIndex)
}

func maxIntWithBits(nbits uint) uint32 {
	// 1 because 2^N bits only gets 2^N-1 if all 1s. +1 more because all 1s is
	// reserved for broadcast.
	const gap float64 = 2

	maxInt := math.Pow(2, float64(nbits))
	if maxInt < gap {
		// we want to avoid underflows, so stick to the point of the API and return
		// effectively zero
		return 0
	}

	return uint32(maxInt - gap)
}

func (s *subnetRequisites) FindSolution() OptimalSubnet {
	opt := OptimalSubnet{}
	_, classCidrOffset, _ := parseip4.Classful(s.ClassfulContext)

	// Brute force solve for Ceil(log2(s.MaxSubnets))
	for {
		if maxIntWithBits(opt.MinSubnetBits) >= uint32(s.MaxSubnets) {
			break
		}
		opt.MinSubnetBits++
	}

	opt.MaxHostsPerSubnet = parseip4.Octets(maxIntWithBits(32 - classCidrOffset - opt.MinSubnetBits))

	mask := parseip4.Octets(0xFFFFFFFF)
	mask <<= (32 - opt.MinSubnetBits)
	opt.Address.Mask = mask.List()

	subnetBitCount := parseip4.CountBitSize(s.SubnetIndex)
	hostBitAddrSpace := 32 - classCidrOffset - subnetBitCount

	ip := s.ClassfulContext.Pack() |
		parseip4.Octets(s.SubnetIndex<<hostBitAddrSpace) |
		s.HostIndex
	opt.Address.IP = ip.List()

	return opt
}

func main() {
	fmt.Printf("part 1: analyzing %d hosts ...\n", len(partOneGivens))
	for _, req := range partOneGivens {
		sol := req.FindSolution()
		fmt.Printf(
			"  given: %s\n\tmin # of subnet bits: %d\n\tmax # hosts per subnet: %d\n\taddress: %s\n",
			req.String(),
			sol.MinSubnetBits,
			sol.MaxHostsPerSubnet,
			sol.Address.String())
	}

	fmt.Printf("part 1 SANITY CHECK: analyzing %d hosts from in-class examples...\n", len(partOneSanityCheck))
	for _, req := range partOneSanityCheck {
		sol := req.FindSolution()
		fmt.Printf(
			"  given: %s\n\tmin # of subnet bits: %d\n\tmax # hosts per subnet: %d\n\taddress: %s\n",
			req.String(),
			sol.MinSubnetBits,
			sol.MaxHostsPerSubnet,
			sol.Address.String())
	}

	fmt.Printf("\npart 2: analyzing %d hosts ...\n", len(partTwoHosts))
	for _, addr := range partTwoHosts {
		classMask, _, klass := parseip4.Classful(addr.IP)

		fmt.Printf(
			"  network: %v (class %s masked)\n\t%v\n\tnetwork id:\t%d\n\t subnet id:\t%d\n\t   host id:\t%d\n",
			(addr.IP.Pack() & classMask.Pack()).List(), klass,
			addr.String(),
			addr.NetworkIndex(),
			addr.SubnetIndex(),
			addr.HostIndex())
	}
}
<|endoftext|>


<commit_before>package tunnel

import (
	"fmt"
	"net"
	"sync"

	"github.com/mmatczuk/go-http-tunnel/id"
	"github.com/mmatczuk/go-http-tunnel/log"
)

// RegistryItem holds information about hosts and listeners associated with a
// client.
type RegistryItem struct {
	Hosts     []*HostAuth
	Listeners []net.Listener
}

// HostAuth holds host and authentication info.
type HostAuth struct {
	Host string
	Auth *Auth
}

type hostInfo struct {
	identifier id.ID
	auth       *Auth
}

type registry struct {
	items  map[id.ID]*RegistryItem
	hosts  map[string]*hostInfo
	mu     sync.RWMutex
	logger log.Logger
}

func newRegistry(logger log.Logger) *registry {
	if logger == nil {
		logger = log.NewNopLogger()
	}

	return &registry{
		items:  make(map[id.ID]*RegistryItem),
		hosts:  make(map[string]*hostInfo, 0),
		logger: logger,
	}
}

var voidRegistryItem = &RegistryItem{}

// Subscribe allows to connect client with a given identifier.
func (r *registry) Subscribe(identifier id.ID) {
	r.logger.Log(
		"level", 1,
		"action", "subscribe",
		"identifier", identifier,
	)

	r.mu.Lock()
	defer r.mu.Unlock()

	if _, ok := r.items[identifier]; ok {
		return
	}

	r.items[identifier] = voidRegistryItem
}

// IsSubscribed returns true if client is subscribed.
func (r *registry) IsSubscribed(identifier id.ID) bool {
	r.mu.RLock()
	defer r.mu.RUnlock()
	_, ok := r.items[identifier]
	return ok
}

// Subscriber returns client identifier assigned to given host.
func (r *registry) Subscriber(hostPort string) (id.ID, *Auth, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	h, ok := r.hosts[trimPort(hostPort)]
	if !ok {
		return id.ID{}, nil, false
	}

	return h.identifier, h.auth, ok
}

// Unsubscribe removes client from registry and returns it's RegistryItem.
func (r *registry) Unsubscribe(identifier id.ID) *RegistryItem {
	r.logger.Log(
		"level", 1,
		"action", "unsubscribe",
		"identifier", identifier,
	)

	r.mu.Lock()
	defer r.mu.Unlock()

	i, ok := r.items[identifier]
	if !ok {
		return nil
	}
	if i == voidRegistryItem {
		return nil
	}

	if i.Hosts != nil {
		for _, h := range i.Hosts {
			delete(r.hosts, h.Host)
		}
	}

	delete(r.items, identifier)

	return i
}

func (r *registry) set(i *RegistryItem, identifier id.ID) error {
	r.logger.Log(
		"level", 2,
		"action", "set registry item",
		"identifier", identifier,
	)

	r.mu.Lock()
	defer r.mu.Unlock()

	j, ok := r.items[identifier]
	if !ok {
		return errClientNotSubscribed
	}
	if j != voidRegistryItem {
		return fmt.Errorf("attempt to overwrite registry item")
	}

	if i.Hosts != nil {
		for _, h := range i.Hosts {
			if h.Auth != nil && h.Auth.User == "" {
				return fmt.Errorf("missing auth user")
			}
			if _, ok := r.hosts[trimPort(h.Host)]; ok {
				return fmt.Errorf("host %q is occupied", h.Host)
			}
		}

		for _, h := range i.Hosts {
			r.hosts[trimPort(h.Host)] = &hostInfo{
				identifier: identifier,
				auth:       h.Auth,
			}
		}
	}

	r.items[identifier] = i

	return nil
}

func (r *registry) clear(identifier id.ID) *RegistryItem {
	r.logger.Log(
		"level", 2,
		"action", "clear registry item",
		"identifier", identifier,
	)

	r.mu.Lock()
	defer r.mu.Unlock()

	i, ok := r.items[identifier]
	if !ok || i == voidRegistryItem {
		return nil
	}

	if i.Hosts != nil {
		for _, h := range i.Hosts {
			delete(r.hosts, trimPort(h.Host))
		}
	}

	r.items[identifier] = voidRegistryItem

	return i
}

func trimPort(hostPort string) (host string) {
	host, _, _ = net.SplitHostPort(hostPort)
	if host == "" {
		host = hostPort
	}
	return
}
<commit_msg>registry: unsubscribe with voidRegisrtyItem fix<commit_after>package tunnel

import (
	"fmt"
	"net"
	"sync"

	"github.com/mmatczuk/go-http-tunnel/id"
	"github.com/mmatczuk/go-http-tunnel/log"
)

// RegistryItem holds information about hosts and listeners associated with a
// client.
type RegistryItem struct {
	Hosts     []*HostAuth
	Listeners []net.Listener
}

// HostAuth holds host and authentication info.
type HostAuth struct {
	Host string
	Auth *Auth
}

type hostInfo struct {
	identifier id.ID
	auth       *Auth
}

type registry struct {
	items  map[id.ID]*RegistryItem
	hosts  map[string]*hostInfo
	mu     sync.RWMutex
	logger log.Logger
}

func newRegistry(logger log.Logger) *registry {
	if logger == nil {
		logger = log.NewNopLogger()
	}

	return &registry{
		items:  make(map[id.ID]*RegistryItem),
		hosts:  make(map[string]*hostInfo, 0),
		logger: logger,
	}
}

var voidRegistryItem = &RegistryItem{}

// Subscribe allows to connect client with a given identifier.
func (r *registry) Subscribe(identifier id.ID) {
	r.logger.Log(
		"level", 1,
		"action", "subscribe",
		"identifier", identifier,
	)

	r.mu.Lock()
	defer r.mu.Unlock()

	if _, ok := r.items[identifier]; ok {
		return
	}

	r.items[identifier] = voidRegistryItem
}

// IsSubscribed returns true if client is subscribed.
func (r *registry) IsSubscribed(identifier id.ID) bool {
	r.mu.RLock()
	defer r.mu.RUnlock()
	_, ok := r.items[identifier]
	return ok
}

// Subscriber returns client identifier assigned to given host.
func (r *registry) Subscriber(hostPort string) (id.ID, *Auth, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	h, ok := r.hosts[trimPort(hostPort)]
	if !ok {
		return id.ID{}, nil, false
	}

	return h.identifier, h.auth, ok
}

// Unsubscribe removes client from registry and returns it's RegistryItem.
func (r *registry) Unsubscribe(identifier id.ID) *RegistryItem {
	r.logger.Log(
		"level", 1,
		"action", "unsubscribe",
		"identifier", identifier,
	)

	r.mu.Lock()
	defer r.mu.Unlock()

	i, ok := r.items[identifier]
	if !ok {
		return nil
	}

	if i.Hosts != nil {
		for _, h := range i.Hosts {
			delete(r.hosts, h.Host)
		}
	}

	delete(r.items, identifier)

	return i
}

func (r *registry) set(i *RegistryItem, identifier id.ID) error {
	r.logger.Log(
		"level", 2,
		"action", "set registry item",
		"identifier", identifier,
	)

	r.mu.Lock()
	defer r.mu.Unlock()

	j, ok := r.items[identifier]
	if !ok {
		return errClientNotSubscribed
	}
	if j != voidRegistryItem {
		return fmt.Errorf("attempt to overwrite registry item")
	}

	if i.Hosts != nil {
		for _, h := range i.Hosts {
			if h.Auth != nil && h.Auth.User == "" {
				return fmt.Errorf("missing auth user")
			}
			if _, ok := r.hosts[trimPort(h.Host)]; ok {
				return fmt.Errorf("host %q is occupied", h.Host)
			}
		}

		for _, h := range i.Hosts {
			r.hosts[trimPort(h.Host)] = &hostInfo{
				identifier: identifier,
				auth:       h.Auth,
			}
		}
	}

	r.items[identifier] = i

	return nil
}

func (r *registry) clear(identifier id.ID) *RegistryItem {
	r.logger.Log(
		"level", 2,
		"action", "clear registry item",
		"identifier", identifier,
	)

	r.mu.Lock()
	defer r.mu.Unlock()

	i, ok := r.items[identifier]
	if !ok || i == voidRegistryItem {
		return nil
	}

	if i.Hosts != nil {
		for _, h := range i.Hosts {
			delete(r.hosts, trimPort(h.Host))
		}
	}

	r.items[identifier] = voidRegistryItem

	return i
}

func trimPort(hostPort string) (host string) {
	host, _, _ = net.SplitHostPort(hostPort)
	if host == "" {
		host = hostPort
	}
	return
}
<|endoftext|>