<commit_before>// Copyright 2020 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
//
// 	https://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 safehttp

import (
	"errors"
	"net/http"
	"net/textproto"
)

// Header represents the key-value pairs in an HTTP header.
// The keys will be in canonical form, as returned by
// textproto.CanonicalMIMEHeaderKey.
type Header struct {
	wrapped http.Header
	claimed map[string]bool
}

func newHeader(h http.Header) Header {
	return Header{wrapped: h, claimed: map[string]bool{}}
}

// Claim claims the header with the given name and returns a function
// which can be used to set the header. The name is first canonicalized
// using textproto.CanonicalMIMEHeaderKey. Other methods in
// the struct can't write to, change or delete the header with this
// name. These methods will instead fail when applied on an claimed
// header. The only way to modify the header is to use the returned
// function. The Set-Cookie header can't be claimed.
func (h Header) Claim(name string) (set func([]string), err error) {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return nil, err
	}
	h.claimed[name] = true
	return func(v []string) {
		h.wrapped[name] = v
	}, nil
}

// Set sets the header with the given name to the given value.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// This method first removes all other values associated with this
// header before setting the new value. Returns an error when
// applied on claimed headers or on the Set-Cookie header.
func (h Header) Set(name, value string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Set(name, value)
	return nil
}

// Add adds a new header with the given name and the given value to
// the collection of headers. The name is first canonicalized using
// textproto.CanonicalMIMEHeaderKey. Returns an error when applied
// on claimed headers or on the Set-Cookie header.
func (h Header) Add(name, value string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Add(name, value)
	return nil
}

// Del deletes all headers with the given name. The name is first
// canonicalized using textproto.CanonicalMIMEHeaderKey. Returns an
// error when applied on claimed headers or on the Set-Cookie
// header.
func (h Header) Del(name string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Del(name)
	return nil
}

// Get returns the value of the first header with the given name.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// If no header exists with the given name then "" is returned.
func (h Header) Get(name string) string {
	return h.wrapped.Get(name)
}

// Values returns all the values of all the headers with the given name.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// The values are returned in the same order as they were sent in the request.
// The values are returned as a copy of the original slice of strings in
// the internal header map. This is to prevent modification of the original
// slice. If no header exists with the given name then an empty slice is
// returned.
func (h Header) Values(name string) []string {
	v := h.wrapped.Values(name)
	clone := make([]string, len(v))
	copy(clone, v)
	return clone
}

// SetCookie adds the cookie provided as a Set-Cookie header in the header
// collection. If the cookie is nil or cookie.Name is invalid, no header is
// added. This is the only method that can modify the Set-Cookie header.
// If other methods try to modify the header they will return errors.
// TODO: Replace http.Cookie with safehttp.Cookie.
func (h Header) SetCookie(cookie *http.Cookie) {
	if v := cookie.String(); v != "" {
		h.wrapped.Add("Set-Cookie", v)
	}
}

// TODO: Add Write, WriteSubset and Clone when needed.

// writableHeader assumes that the given name already has been canonicalized
// using textproto.CanonicalMIMEHeaderKey.
func (h Header) writableHeader(name string) error {
	// TODO(@mattiasgrenfeldt, @kele, @empijei): Think about how this should
	// work during legacy conversions.
	if name == "Set-Cookie" {
		return errors.New("can't write to Set-Cookie header")
	}
	if h.claimed[name] {
		return errors.New("claimed header")
	}
	return nil
}
<commit_msg>Split construction of struct to multiple lines<commit_after>// Copyright 2020 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
//
// 	https://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 safehttp

import (
	"errors"
	"net/http"
	"net/textproto"
)

// Header represents the key-value pairs in an HTTP header.
// The keys will be in canonical form, as returned by
// textproto.CanonicalMIMEHeaderKey.
type Header struct {
	wrapped http.Header
	claimed map[string]bool
}

func newHeader(h http.Header) Header {
	return Header{
		wrapped: h,
		claimed: map[string]bool{},
	}
}

// Claim claims the header with the given name and returns a function
// which can be used to set the header. The name is first canonicalized
// using textproto.CanonicalMIMEHeaderKey. Other methods in
// the struct can't write to, change or delete the header with this
// name. These methods will instead fail when applied on an claimed
// header. The only way to modify the header is to use the returned
// function. The Set-Cookie header can't be claimed.
func (h Header) Claim(name string) (set func([]string), err error) {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return nil, err
	}
	h.claimed[name] = true
	return func(v []string) {
		h.wrapped[name] = v
	}, nil
}

// Set sets the header with the given name to the given value.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// This method first removes all other values associated with this
// header before setting the new value. Returns an error when
// applied on claimed headers or on the Set-Cookie header.
func (h Header) Set(name, value string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Set(name, value)
	return nil
}

// Add adds a new header with the given name and the given value to
// the collection of headers. The name is first canonicalized using
// textproto.CanonicalMIMEHeaderKey. Returns an error when applied
// on claimed headers or on the Set-Cookie header.
func (h Header) Add(name, value string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Add(name, value)
	return nil
}

// Del deletes all headers with the given name. The name is first
// canonicalized using textproto.CanonicalMIMEHeaderKey. Returns an
// error when applied on claimed headers or on the Set-Cookie
// header.
func (h Header) Del(name string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Del(name)
	return nil
}

// Get returns the value of the first header with the given name.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// If no header exists with the given name then "" is returned.
func (h Header) Get(name string) string {
	return h.wrapped.Get(name)
}

// Values returns all the values of all the headers with the given name.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// The values are returned in the same order as they were sent in the request.
// The values are returned as a copy of the original slice of strings in
// the internal header map. This is to prevent modification of the original
// slice. If no header exists with the given name then an empty slice is
// returned.
func (h Header) Values(name string) []string {
	v := h.wrapped.Values(name)
	clone := make([]string, len(v))
	copy(clone, v)
	return clone
}

// SetCookie adds the cookie provided as a Set-Cookie header in the header
// collection. If the cookie is nil or cookie.Name is invalid, no header is
// added. This is the only method that can modify the Set-Cookie header.
// If other methods try to modify the header they will return errors.
// TODO: Replace http.Cookie with safehttp.Cookie.
func (h Header) SetCookie(cookie *http.Cookie) {
	if v := cookie.String(); v != "" {
		h.wrapped.Add("Set-Cookie", v)
	}
}

// TODO: Add Write, WriteSubset and Clone when needed.

// writableHeader assumes that the given name already has been canonicalized
// using textproto.CanonicalMIMEHeaderKey.
func (h Header) writableHeader(name string) error {
	// TODO(@mattiasgrenfeldt, @kele, @empijei): Think about how this should
	// work during legacy conversions.
	if name == "Set-Cookie" {
		return errors.New("can't write to Set-Cookie header")
	}
	if h.claimed[name] {
		return errors.New("claimed header")
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 * Minimal object storage library (C) 2015 Minio, Inc.
 *
 * 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 minio

// bucketHandler is an http.Handler that verifies bucket responses and validates incoming requests
import (
	"bytes"
	"io"
	"net/http"
	"strconv"
	"time"
)

type bucketHandler struct {
	resource string
}

func (h bucketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		switch {
		case r.URL.Path == "/":
			response := []byte("<ListAllMyBucketsResult xmlns=\"http://doc.s3.amazonaws.com/2006-03-01\"><Buckets><Bucket><Name>bucket</Name><CreationDate>2015-05-20T23:05:09.230Z</CreationDate></Bucket></Buckets><Owner><ID>minio</ID><DisplayName>minio</DisplayName></Owner></ListAllMyBucketsResult>")
			w.Header().Set("Content-Length", strconv.Itoa(len(response)))
			w.Write(response)
		case r.URL.Path == "/bucket":
			_, ok := r.URL.Query()["acl"]
			if ok {
				response := []byte("<AccessControlPolicy><Owner><ID>75aa57f09aa0c8caeab4f8c24e99d10f8e7faeebf76c078efc7c6caea54ba06a</ID><DisplayName>CustomersName@amazon.com</DisplayName></Owner><AccessControlList><Grant><Grantee xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\" xsi:type=\"CanonicalUser\"><ID>75aa57f09aa0c8caeab4f8c24e99d10f8e7faeebf76c078efc7c6caea54ba06a</ID><DisplayName>CustomersName@amazon.com</DisplayName></Grantee><Permission>FULL_CONTROL</Permission></Grant></AccessControlList></AccessControlPolicy>")
				w.Header().Set("Content-Length", strconv.Itoa(len(response)))
				w.Write(response)
				return
			}
			fallthrough
		case r.URL.Path == "/bucket":
			response := []byte("<ListBucketResult xmlns=\"http://doc.s3.amazonaws.com/2006-03-01\"><Contents><ETag>259d04a13802ae09c7e41be50ccc6baa</ETag><Key>object</Key><LastModified>2015-05-21T18:24:21.097Z</LastModified><Size>22061</Size><Owner><ID>minio</ID><DisplayName>minio</DisplayName></Owner><StorageClass>STANDARD</StorageClass></Contents><Delimiter></Delimiter><EncodingType></EncodingType><IsTruncated>false</IsTruncated><Marker></Marker><MaxKeys>1000</MaxKeys><Name>testbucket</Name><NextMarker></NextMarker><Prefix></Prefix></ListBucketResult>")
			w.Header().Set("Content-Length", strconv.Itoa(len(response)))
			w.Write(response)
		}
	case r.Method == "PUT":
		switch {
		case r.URL.Path == h.resource:
			_, ok := r.URL.Query()["acl"]
			if ok {
				if r.Header.Get("x-amz-acl") != "public-read-write" {
					w.WriteHeader(http.StatusNotImplemented)
				}
			}
			w.WriteHeader(http.StatusOK)
		default:
			w.WriteHeader(http.StatusBadRequest)
		}
	case r.Method == "HEAD":
		switch {
		case r.URL.Path == h.resource:
			w.WriteHeader(http.StatusOK)
		default:
			w.WriteHeader(http.StatusForbidden)
		}
	case r.Method == "DELETE":
		switch {
		case r.URL.Path != h.resource:
			w.WriteHeader(http.StatusNotFound)
		default:
			h.resource = ""
			w.WriteHeader(http.StatusOK)
		}
	}
}

// objectHandler is an http.Handler that verifies object responses and validates incoming requests
type objectHandler struct {
	resource string
	data     []byte
}

func (h objectHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "PUT":
		length, err := strconv.Atoi(r.Header.Get("Content-Length"))
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		var buffer bytes.Buffer
		_, err = io.CopyN(&buffer, r.Body, int64(length))
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		if !bytes.Equal(h.data, buffer.Bytes()) {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		w.Header().Set("ETag", "9af2f8218b150c351ad802c6f3d66abe")
		w.WriteHeader(http.StatusOK)
	case r.Method == "HEAD":
		if r.URL.Path != h.resource {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		w.Header().Set("Content-Length", strconv.Itoa(len(h.data)))
		w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
		w.Header().Set("ETag", "9af2f8218b150c351ad802c6f3d66abe")
		w.WriteHeader(http.StatusOK)
	case r.Method == "POST":
	case r.Method == "GET":
		if r.URL.Path != h.resource {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		w.Header().Set("Content-Length", strconv.Itoa(len(h.data)))
		w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
		w.Header().Set("ETag", "9af2f8218b150c351ad802c6f3d66abe")
		w.WriteHeader(http.StatusOK)
		io.Copy(w, bytes.NewReader(h.data))
	case r.Method == "DELETE":
		if r.URL.Path != h.resource {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		h.resource = ""
		h.data = nil
		w.WriteHeader(http.StatusOK)
	}
}
<commit_msg>Update test handlers to use new XML<commit_after>/*
 * Minimal object storage library (C) 2015 Minio, Inc.
 *
 * 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 minio

// bucketHandler is an http.Handler that verifies bucket responses and validates incoming requests
import (
	"bytes"
	"io"
	"net/http"
	"strconv"
	"time"
)

type bucketHandler struct {
	resource string
}

func (h bucketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		switch {
		case r.URL.Path == "/":
			response := []byte("<?xml version=\"1.0\" encoding=\"UTF-8\"?><ListAllMyBucketsResult xmlns=\"http://doc.s3.amazonaws.com/2006-03-01\"><Buckets><Bucket><Name>bucket</Name><CreationDate>2015-05-20T23:05:09.230Z</CreationDate></Bucket></Buckets><Owner><ID>minio</ID><DisplayName>minio</DisplayName></Owner></ListAllMyBucketsResult>")
			w.Header().Set("Content-Length", strconv.Itoa(len(response)))
			w.Write(response)
		case r.URL.Path == "/bucket":
			_, ok := r.URL.Query()["acl"]
			if ok {
				response := []byte("<?xml version=\"1.0\" encoding=\"UTF-8\"?><AccessControlPolicy><Owner><ID>75aa57f09aa0c8caeab4f8c24e99d10f8e7faeebf76c078efc7c6caea54ba06a</ID><DisplayName>CustomersName@amazon.com</DisplayName></Owner><AccessControlList><Grant><Grantee xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\" xsi:type=\"CanonicalUser\"><ID>75aa57f09aa0c8caeab4f8c24e99d10f8e7faeebf76c078efc7c6caea54ba06a</ID><DisplayName>CustomersName@amazon.com</DisplayName></Grantee><Permission>FULL_CONTROL</Permission></Grant></AccessControlList></AccessControlPolicy>")
				w.Header().Set("Content-Length", strconv.Itoa(len(response)))
				w.Write(response)
				return
			}
			fallthrough
		case r.URL.Path == "/bucket":
			response := []byte("<?xml version=\"1.0\" encoding=\"UTF-8\"?><ListBucketResult xmlns=\"http://doc.s3.amazonaws.com/2006-03-01\"><Contents><ETag>259d04a13802ae09c7e41be50ccc6baa</ETag><Key>object</Key><LastModified>2015-05-21T18:24:21.097Z</LastModified><Size>22061</Size><Owner><ID>minio</ID><DisplayName>minio</DisplayName></Owner><StorageClass>STANDARD</StorageClass></Contents><Delimiter></Delimiter><EncodingType></EncodingType><IsTruncated>false</IsTruncated><Marker></Marker><MaxKeys>1000</MaxKeys><Name>testbucket</Name><NextMarker></NextMarker><Prefix></Prefix></ListBucketResult>")
			w.Header().Set("Content-Length", strconv.Itoa(len(response)))
			w.Write(response)
		}
	case r.Method == "PUT":
		switch {
		case r.URL.Path == h.resource:
			_, ok := r.URL.Query()["acl"]
			if ok {
				if r.Header.Get("x-amz-acl") != "public-read-write" {
					w.WriteHeader(http.StatusNotImplemented)
				}
			}
			w.WriteHeader(http.StatusOK)
		default:
			w.WriteHeader(http.StatusBadRequest)
		}
	case r.Method == "HEAD":
		switch {
		case r.URL.Path == h.resource:
			w.WriteHeader(http.StatusOK)
		default:
			w.WriteHeader(http.StatusForbidden)
		}
	case r.Method == "DELETE":
		switch {
		case r.URL.Path != h.resource:
			w.WriteHeader(http.StatusNotFound)
		default:
			h.resource = ""
			w.WriteHeader(http.StatusOK)
		}
	}
}

// objectHandler is an http.Handler that verifies object responses and validates incoming requests
type objectHandler struct {
	resource string
	data     []byte
}

func (h objectHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "PUT":
		length, err := strconv.Atoi(r.Header.Get("Content-Length"))
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		var buffer bytes.Buffer
		_, err = io.CopyN(&buffer, r.Body, int64(length))
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		if !bytes.Equal(h.data, buffer.Bytes()) {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		w.Header().Set("ETag", "9af2f8218b150c351ad802c6f3d66abe")
		w.WriteHeader(http.StatusOK)
	case r.Method == "HEAD":
		if r.URL.Path != h.resource {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		w.Header().Set("Content-Length", strconv.Itoa(len(h.data)))
		w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
		w.Header().Set("ETag", "9af2f8218b150c351ad802c6f3d66abe")
		w.WriteHeader(http.StatusOK)
	case r.Method == "POST":
		_, ok := r.URL.Query()["uploads"]
		if ok {
			response := []byte("<?xml version=\"1.0\" encoding=\"UTF-8\"?><InitiateMultipartUploadResult xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\"><Bucket>example-bucket</Bucket><Key>object</Key><UploadId>XXBsb2FkIElEIGZvciBlbHZpbmcncyVcdS1tb3ZpZS5tMnRzEEEwbG9hZA</UploadId></InitiateMultipartUploadResult>")
			w.Header().Set("Content-Length", strconv.Itoa(len(response)))
			w.Write(response)
			return
		}
	case r.Method == "GET":
		_, ok := r.URL.Query()["uploadId"]
		if ok {
			uploadID := r.URL.Query().Get("uploadId")
			if uploadID != "XXBsb2FkIElEIGZvciBlbHZpbmcncyVcdS1tb3ZpZS5tMnRzEEEwbG9hZA" {
				w.WriteHeader(http.StatusNotFound)
				return
			}
			response := []byte("<?xml version=\"1.0\" encoding=\"UTF-8\"?><ListPartsResult xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\"><Bucket>example-bucket</Bucket><Key>example-object</Key><UploadId>XXBsb2FkIElEIGZvciBlbHZpbmcncyVcdS1tb3ZpZS5tMnRzEEEwbG9hZA</UploadId><Initiator><ID>arn:aws:iam::111122223333:user/some-user-11116a31-17b5-4fb7-9df5-b288870f11xx</ID><DisplayName>umat-user-11116a31-17b5-4fb7-9df5-b288870f11xx</DisplayName></Initiator><Owner><ID>75aa57f09aa0c8caeab4f8c24e99d10f8e7faeebf76c078efc7c6caea54ba06a</ID><DisplayName>someName</DisplayName></Owner><StorageClass>STANDARD</StorageClass><PartNumberMarker>1</PartNumberMarker><NextPartNumberMarker>3</NextPartNumberMarker><MaxParts>2</MaxParts><IsTruncated>true</IsTruncated><Part><PartNumber>2</PartNumber><LastModified>2010-11-10T20:48:34.000Z</LastModified><ETag>\"7778aef83f66abc1fa1e8477f296d394\"</ETag><Size>10485760</Size></Part><Part><PartNumber>3</PartNumber><LastModified>2010-11-10T20:48:33.000Z</LastModified><ETag>\"aaaa18db4cc2f85cedef654fccc4a4x8\"</ETag><Size>10485760</Size></Part></ListPartsResult>")
			w.Header().Set("Content-Length", strconv.Itoa(len(response)))
			w.Write(response)
			return
		}
		if r.URL.Path != h.resource {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		w.Header().Set("Content-Length", strconv.Itoa(len(h.data)))
		w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
		w.Header().Set("ETag", "9af2f8218b150c351ad802c6f3d66abe")
		w.WriteHeader(http.StatusOK)
		io.Copy(w, bytes.NewReader(h.data))
	case r.Method == "DELETE":
		if r.URL.Path != h.resource {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		h.resource = ""
		h.data = nil
		w.WriteHeader(http.StatusOK)
	}
}
<|endoftext|>


<commit_before>package tests

import (
	"github.com/bozaro/tech-db-forum/generated/client"
	"github.com/bozaro/tech-db-forum/generated/client/operations"
	"github.com/bozaro/tech-db-forum/generated/models"
	"github.com/go-openapi/strfmt"
	"sort"
	"time"
)

func init() {
	Register(Checker{
		Name:        "forum_get_threads_simple",
		Description: "",
		FnCheck:     Modifications(CheckForumGetThreadsSimple),
		Deps: []string{
			"thread_create_simple",
		},
	})
	Register(Checker{
		Name:        "forum_get_threads_notfound",
		Description: "",
		FnCheck:     Modifications(CheckForumGetThreadsNotFound),
		Deps: []string{
			"thread_create_simple",
		},
	})
}

func CheckForumGetThreadsSimple(c *client.Forum, m *Modify) {
	forum := CreateForum(c, nil, nil)
	threads := []models.Thread{}
	created := time.Now()
	created.Round(time.Millisecond)
	for i := 0; i < 10; i++ {
		thread := CreateThread(c, nil, forum, nil)
		threads = append(threads, *thread)
	}
	sort.Sort(ThreadByCreated(threads))

	var desc *bool

	// Desc
	small := time.Millisecond
	switch m.Int(3) {
	case 1:
		v := bool(true)
		small = -small
		desc = &v
		sort.Sort(sort.Reverse(ThreadByCreated(threads)))
	case 2:
		v := bool(false)
		desc = &v
	}

	// Check read all
	c.Operations.ForumGetThreads(operations.NewForumGetThreadsParams().
		WithSlug(forum.Slug).
		WithDesc(desc).
		WithContext(Expected(200, &threads, nil)))

	// Check read by 4 records
	limit := int32(4)
	var since *strfmt.DateTime = nil
	for n := 0; n < len(threads); n += int(limit) - 1 {
		m := n + int(limit)
		if m > len(threads) {
			m = len(threads)
		}
		expected := threads[n:m]
		c.Operations.ForumGetThreads(operations.NewForumGetThreadsParams().
			WithSlug(forum.Slug).
			WithLimit(&limit).
			WithDesc(desc).
			WithSince(since).
			WithContext(Expected(200, &expected, nil)))
		since = threads[m-1].Created
	}

	// Check read after all
	after_last := strfmt.DateTime(time.Time(*threads[len(threads)-1].Created).Add(small))
	c.Operations.ForumGetThreads(operations.NewForumGetThreadsParams().
		WithSlug(forum.Slug).
		WithLimit(&limit).
		WithDesc(desc).
		WithSince(&after_last).
		WithContext(Expected(200, &[]models.Thread{}, nil)))
}

func CheckForumGetThreadsNotFound(c *client.Forum, m *Modify) {
	var limit *int32
	var since *strfmt.DateTime
	var desc *bool

	forum := RandomForum()
	// Limit
	if m.Bool() {
		v := int32(10)
		limit = &v
	}
	// Since
	if m.Bool() {
		v := strfmt.DateTime(time.Now())
		since = &v
	}
	// Desc
	switch m.Int(3) {
	case 1:
		v := bool(true)
		desc = &v
	case 2:
		v := bool(false)
		desc = &v
	}

	// Check
	_, err := c.Operations.ForumGetThreads(operations.NewForumGetThreadsParams().
		WithSlug(forum.Slug).
		WithLimit(limit).
		WithSince(since).
		WithDesc(desc).
		WithContext(Expected(404, nil, nil)))
	CheckIsType(operations.NewForumGetThreadsNotFound(), err)
}
<commit_msg>Поправил проверку дат для того, чтобы не зависеть от варианта передачи временной зоны<commit_after>package tests

import (
	"github.com/bozaro/tech-db-forum/generated/client"
	"github.com/bozaro/tech-db-forum/generated/client/operations"
	"github.com/bozaro/tech-db-forum/generated/models"
	"github.com/go-openapi/strfmt"
	"sort"
	"time"
)

func init() {
	Register(Checker{
		Name:        "forum_get_threads_simple",
		Description: "",
		FnCheck:     Modifications(CheckForumGetThreadsSimple),
		Deps: []string{
			"thread_create_simple",
		},
	})
	Register(Checker{
		Name:        "forum_get_threads_notfound",
		Description: "",
		FnCheck:     Modifications(CheckForumGetThreadsNotFound),
		Deps: []string{
			"thread_create_simple",
		},
	})
}

func filterThreads(data interface{}) interface{} {
	threads := data.(*[]models.Thread)
	for i := range *threads {
		thread := &(*threads)[i]
		if thread.Created != nil {
			created := strfmt.DateTime(time.Time(*thread.Created).UTC())
			thread.Created = &created
		}
	}
	return threads
}

func CheckForumGetThreadsSimple(c *client.Forum, m *Modify) {
	forum := CreateForum(c, nil, nil)
	threads := []models.Thread{}
	created := time.Now()
	created.Round(time.Millisecond)
	for i := 0; i < 10; i++ {
		thread := CreateThread(c, nil, forum, nil)
		threads = append(threads, *thread)
	}
	sort.Sort(ThreadByCreated(threads))

	var desc *bool

	// Desc
	small := time.Millisecond
	switch m.Int(3) {
	case 1:
		v := bool(true)
		small = -small
		desc = &v
		sort.Sort(sort.Reverse(ThreadByCreated(threads)))
	case 2:
		v := bool(false)
		desc = &v
	}

	// Check read all
	c.Operations.ForumGetThreads(operations.NewForumGetThreadsParams().
		WithSlug(forum.Slug).
		WithDesc(desc).
		WithContext(Expected(200, &threads, filterThreads)))

	// Check read by 4 records
	limit := int32(4)
	var since *strfmt.DateTime = nil
	for n := 0; n < len(threads); n += int(limit) - 1 {
		m := n + int(limit)
		if m > len(threads) {
			m = len(threads)
		}
		expected := threads[n:m]
		c.Operations.ForumGetThreads(operations.NewForumGetThreadsParams().
			WithSlug(forum.Slug).
			WithLimit(&limit).
			WithDesc(desc).
			WithSince(since).
			WithContext(Expected(200, &expected, filterThreads)))
		since = threads[m-1].Created
	}

	// Check read after all
	after_last := strfmt.DateTime(time.Time(*threads[len(threads)-1].Created).Add(small))
	c.Operations.ForumGetThreads(operations.NewForumGetThreadsParams().
		WithSlug(forum.Slug).
		WithLimit(&limit).
		WithDesc(desc).
		WithSince(&after_last).
		WithContext(Expected(200, &[]models.Thread{}, nil)))
}

func CheckForumGetThreadsNotFound(c *client.Forum, m *Modify) {
	var limit *int32
	var since *strfmt.DateTime
	var desc *bool

	forum := RandomForum()
	// Limit
	if m.Bool() {
		v := int32(10)
		limit = &v
	}
	// Since
	if m.Bool() {
		v := strfmt.DateTime(time.Now())
		since = &v
	}
	// Desc
	switch m.Int(3) {
	case 1:
		v := bool(true)
		desc = &v
	case 2:
		v := bool(false)
		desc = &v
	}

	// Check
	_, err := c.Operations.ForumGetThreads(operations.NewForumGetThreadsParams().
		WithSlug(forum.Slug).
		WithLimit(limit).
		WithSince(since).
		WithDesc(desc).
		WithContext(Expected(404, nil, nil)))
	CheckIsType(operations.NewForumGetThreadsNotFound(), err)
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2016 Intel Corporation
//
// 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 main

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"

	"github.com/ciao-project/ciao/ciao-controller/api"
	"github.com/ciao-project/ciao/ciao-controller/types"
	"github.com/ciao-project/ciao/openstack/compute"
	"github.com/ciao-project/ciao/payloads"
	"github.com/intel/tfortools"

	"gopkg.in/yaml.v2"
)

var workloadCommand = &command{
	SubCommands: map[string]subCommand{
		"list":   new(workloadListCommand),
		"create": new(workloadCreateCommand),
		"delete": new(workloadDeleteCommand),
		"show":   new(workloadShowCommand),
	},
}

type workloadListCommand struct {
	Flag     flag.FlagSet
	template string
}

func (cmd *workloadListCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload list

List all workloads

`)
	cmd.Flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n%s",
		tfortools.GenerateUsageDecorated("f", compute.FlavorsDetails{}.Flavors, nil))
	os.Exit(2)
}

func (cmd *workloadListCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.template, "f", "", "Template used to format output")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadListCommand) run(args []string) error {
	if *tenantID == "" {
		fatalf("Missing required -tenant-id parameter")
	}

	var flavors compute.FlavorsDetails
	if *tenantID == "" {
		*tenantID = "faketenant"
	}

	url := buildComputeURL("%s/flavors/detail", *tenantID)

	resp, err := sendHTTPRequest("GET", url, nil, nil)
	if err != nil {
		fatalf(err.Error())
	}

	err = unmarshalHTTPResponse(resp, &flavors)
	if err != nil {
		fatalf(err.Error())
	}

	if cmd.template != "" {
		return tfortools.OutputToTemplate(os.Stdout, "workload-list", cmd.template,
			&flavors.Flavors, nil)
	}

	for i, flavor := range flavors.Flavors {
		fmt.Printf("Workload %d\n", i+1)
		fmt.Printf("\tName: %s\n\tUUID:%s\n\tCPUs: %d\n\tMemory: %d MB\n",
			flavor.Name, flavor.ID, flavor.Vcpus, flavor.RAM)
	}
	return nil
}

type workloadCreateCommand struct {
	Flag     flag.FlagSet
	yamlFile string
}

func (cmd *workloadCreateCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.yamlFile, "yaml", "", "filename for yaml which describes the workload")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadCreateCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload create [flags]

Create a new workload

The create flags are:

`)
	cmd.Flag.PrintDefaults()
	os.Exit(2)
}

func getCiaoWorkloadsResource() (string, error) {
	return getCiaoResource("workloads", api.WorkloadsV1)
}

type source struct {
	Type types.SourceType `yaml:"service"`
	ID   string           `yaml:"id"`
}

type disk struct {
	ID        *string `yaml:"volume_id,omitempty"`
	Size      int     `yaml:"size"`
	Bootable  bool    `yaml:"bootable"`
	Source    source  `yaml:"source"`
	Ephemeral bool    `yaml:"ephemeral"`
}

type defaultResources struct {
	VCPUs int `yaml:"vcpus"`
	MemMB int `yaml:"mem_mb"`
}

// we currently only use the first disk due to lack of support
// in types.Workload for multiple storage resources.
type workloadOptions struct {
	Description     string           `yaml:"description"`
	VMType          string           `yaml:"vm_type"`
	FWType          string           `yaml:"fw_type,omitempty"`
	ImageName       string           `yaml:"image_name,omitempty"`
	Defaults        defaultResources `yaml:"defaults"`
	CloudConfigFile string           `yaml:"cloud_init,omitempty"`
	Disks           []disk           `yaml:"disks,omitempty"`
}

func optToReqStorage(opt workloadOptions) ([]types.StorageResource, error) {
	storage := make([]types.StorageResource, 0)
	bootableCount := 0
	for _, disk := range opt.Disks {
		res := types.StorageResource{
			Size:      disk.Size,
			Bootable:  disk.Bootable,
			Ephemeral: disk.Ephemeral,
		}

		// Use existing volume
		if disk.ID != nil {
			res.ID = *disk.ID
		} else {
			// Create a new one
			if disk.Source.Type == "" {
				disk.Source.Type = types.Empty
			}

			if disk.Source.Type != types.Empty {
				res.SourceType = disk.Source.Type
				res.SourceID = disk.Source.ID

				if res.SourceID == "" {
					return nil, errors.New("Invalid workload yaml: when using a source an id must also be specified")
				}
			} else {
				if disk.Bootable == true {
					// you may not request a bootable drive
					// from an empty source
					return nil, errors.New("Invalid workload yaml: empty disk source may not be bootable")
				}

				if disk.Size <= 0 {
					return nil, errors.New("Invalid workload yaml: size required when creating a volume")
				}
			}
		}

		if disk.Bootable {
			bootableCount++
		}

		storage = append(storage, res)
	}

	if payloads.Hypervisor(opt.VMType) == payloads.QEMU && bootableCount == 0 {
		return nil, errors.New("Invalid workload yaml: no bootable disks specified for a VM")
	}

	return storage, nil
}

func optToReq(opt workloadOptions, req *types.Workload) error {
	b, err := ioutil.ReadFile(opt.CloudConfigFile)
	if err != nil {
		return err
	}

	config := string(b)

	// this is where you'd validate that the options make
	// sense.
	req.Description = opt.Description
	req.VMType = payloads.Hypervisor(opt.VMType)
	req.FWType = opt.FWType
	req.ImageName = opt.ImageName
	req.Config = config
	req.Storage, err = optToReqStorage(opt)

	if err != nil {
		return err
	}

	// all default resources are required.
	defaults := opt.Defaults

	r := payloads.RequestedResource{
		Type:  payloads.VCPUs,
		Value: defaults.VCPUs,
	}
	req.Defaults = append(req.Defaults, r)

	r = payloads.RequestedResource{
		Type:  payloads.MemMB,
		Value: defaults.MemMB,
	}
	req.Defaults = append(req.Defaults, r)

	return nil
}

func outputWorkload(w types.Workload) {
	var opt workloadOptions

	opt.Description = w.Description
	opt.VMType = string(w.VMType)
	opt.FWType = w.FWType
	opt.ImageName = w.ImageName
	for _, d := range w.Defaults {
		if d.Type == payloads.VCPUs {
			opt.Defaults.VCPUs = d.Value
		} else if d.Type == payloads.MemMB {
			opt.Defaults.MemMB = d.Value
		}
	}

	for _, s := range w.Storage {
		d := disk{
			Size:      s.Size,
			Bootable:  s.Bootable,
			Ephemeral: s.Ephemeral,
		}
		if s.ID != "" {
			d.ID = &s.ID
		}

		src := source{
			Type: s.SourceType,
			ID:   s.SourceID,
		}

		d.Source = src

		opt.Disks = append(opt.Disks, d)
	}

	b, err := yaml.Marshal(opt)
	if err != nil {
		fatalf(err.Error())
	}

	fmt.Println(string(b))
	fmt.Println(w.Config)
}

func (cmd *workloadCreateCommand) run(args []string) error {
	var opt workloadOptions
	var req types.Workload

	if cmd.yamlFile == "" {
		cmd.usage()
	}

	f, err := ioutil.ReadFile(cmd.yamlFile)
	if err != nil {
		fatalf("Unable to read workload config file: %s\n", err)
	}

	err = yaml.Unmarshal(f, &opt)
	if err != nil {
		fatalf("Config file invalid: %s\n", err)
	}

	err = optToReq(opt, &req)
	if err != nil {
		fatalf(err.Error())
	}

	b, err := json.Marshal(req)
	if err != nil {
		fatalf(err.Error())
	}

	body := bytes.NewReader(b)

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	resp, err := sendCiaoRequest("POST", url, nil, body, ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusCreated {
		fatalf("Workload creation failed: %s", resp.Status)
	}

	var workload types.WorkloadResponse

	err = unmarshalHTTPResponse(resp, &workload)
	if err != nil {
		fatalf(err.Error())
	}

	fmt.Printf("Created new workload: %s\n", workload.Workload.ID)

	return nil
}

type workloadDeleteCommand struct {
	Flag     flag.FlagSet
	workload string
}

func (cmd *workloadDeleteCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload delete [flags]

Deletes a given workload

The delete flags are:

`)
	cmd.Flag.PrintDefaults()
	os.Exit(2)
}

func (cmd *workloadDeleteCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.workload, "workload", "", "Workload UUID")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadDeleteCommand) run(args []string) error {
	if cmd.workload == "" {
		cmd.usage()
	}

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	// you should do a get first and search for the workload,
	// then use the href - but not with the currently used
	// OpenStack API. Until we support GET with a ciao API,
	// just hard code the path.
	url = fmt.Sprintf("%s/%s", url, cmd.workload)

	resp, err := sendCiaoRequest("DELETE", url, nil, nil, ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusNoContent {
		fatalf("Workload deletion failed: %s", resp.Status)
	}

	return nil
}

type workloadShowCommand struct {
	Flag     flag.FlagSet
	template string
	workload string
}

func (cmd *workloadShowCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload show

Show workload details

`)
	cmd.Flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n%s",
		tfortools.GenerateUsageDecorated("f", types.Workload{}, nil))
	os.Exit(2)
}

func (cmd *workloadShowCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.workload, "workload", "", "Workload UUID")
	cmd.Flag.StringVar(&cmd.template, "f", "", "Template used to format output")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadShowCommand) run(args []string) error {
	var wl types.Workload

	if cmd.workload == "" {
		cmd.usage()
	}

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	// you should do a get first and search for the workload,
	// then use the href - but not with the currently used
	// OpenStack API. Until we support GET with a ciao API,
	// just hard code the path.
	url = fmt.Sprintf("%s/%s", url, cmd.workload)

	resp, err := sendCiaoRequest("GET", url, nil, nil, ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusOK {
		fatalf("Workload show failed: %s", resp.Status)
	}

	err = unmarshalHTTPResponse(resp, &wl)
	if err != nil {
		fatalf(err.Error())
	}

	if cmd.template != "" {
		return tfortools.OutputToTemplate(os.Stdout, "workload-show", cmd.template, &wl, nil)
	}

	outputWorkload(wl)
	return nil
}
<commit_msg>ciao-cli: Use workload API for listing workloads<commit_after>//
// Copyright (c) 2016 Intel Corporation
//
// 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 main

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"

	"github.com/ciao-project/ciao/ciao-controller/api"
	"github.com/ciao-project/ciao/ciao-controller/types"
	"github.com/ciao-project/ciao/payloads"
	"github.com/intel/tfortools"

	"gopkg.in/yaml.v2"
)

var workloadCommand = &command{
	SubCommands: map[string]subCommand{
		"list":   new(workloadListCommand),
		"create": new(workloadCreateCommand),
		"delete": new(workloadDeleteCommand),
		"show":   new(workloadShowCommand),
	},
}

type workloadListCommand struct {
	Flag     flag.FlagSet
	template string
}

// Workload contains detailed information about a workload
type Workload struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	CPUs int    `json:"vcpus"`
	Mem  int    `json:"ram"`
}

func (cmd *workloadListCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload list

List all workloads

`)
	cmd.Flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n%s",
		tfortools.GenerateUsageDecorated("f", []Workload{}, nil))
	os.Exit(2)
}

func (cmd *workloadListCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.template, "f", "", "Template used to format output")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadListCommand) run(args []string) error {
	if *tenantID == "" {
		fatalf("Missing required -tenant-id parameter")
	}

	var wls []types.Workload

	var url string
	if checkPrivilege() {
		url = buildCiaoURL("workloads")
	} else {
		url = buildCiaoURL("%s/workloads", *tenantID)
	}

	resp, err := sendCiaoRequest("GET", url, nil, nil, api.WorkloadsV1)
	if err != nil {
		fatalf(err.Error())
	}

	err = unmarshalHTTPResponse(resp, &wls)
	if err != nil {
		fatalf(err.Error())
	}

	var workloads []Workload
	for i, wl := range wls {
		workloads = append(workloads, Workload{
			Name: wl.Description,
			ID:   wl.ID,
		})

		for _, r := range wl.Defaults {
			if r.Type == payloads.MemMB {
				workloads[i].Mem = r.Value
			}
			if r.Type == payloads.VCPUs {
				workloads[i].CPUs = r.Value
			}
		}
	}

	if cmd.template != "" {
		return tfortools.OutputToTemplate(os.Stdout, "workload-list", cmd.template,
			workloads, nil)
	}

	for i, wl := range workloads {
		fmt.Printf("Workload %d\n", i+1)
		fmt.Printf("\tName: %s\n\tUUID:%s\n\tCPUs: %d\n\tMemory: %d MB\n",
			wl.Name, wl.ID, wl.CPUs, wl.Mem)
	}

	return nil
}

type workloadCreateCommand struct {
	Flag     flag.FlagSet
	yamlFile string
}

func (cmd *workloadCreateCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.yamlFile, "yaml", "", "filename for yaml which describes the workload")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadCreateCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload create [flags]

Create a new workload

The create flags are:

`)
	cmd.Flag.PrintDefaults()
	os.Exit(2)
}

func getCiaoWorkloadsResource() (string, error) {
	return getCiaoResource("workloads", api.WorkloadsV1)
}

type source struct {
	Type types.SourceType `yaml:"service"`
	ID   string           `yaml:"id"`
}

type disk struct {
	ID        *string `yaml:"volume_id,omitempty"`
	Size      int     `yaml:"size"`
	Bootable  bool    `yaml:"bootable"`
	Source    source  `yaml:"source"`
	Ephemeral bool    `yaml:"ephemeral"`
}

type defaultResources struct {
	VCPUs int `yaml:"vcpus"`
	MemMB int `yaml:"mem_mb"`
}

// we currently only use the first disk due to lack of support
// in types.Workload for multiple storage resources.
type workloadOptions struct {
	Description     string           `yaml:"description"`
	VMType          string           `yaml:"vm_type"`
	FWType          string           `yaml:"fw_type,omitempty"`
	ImageName       string           `yaml:"image_name,omitempty"`
	Defaults        defaultResources `yaml:"defaults"`
	CloudConfigFile string           `yaml:"cloud_init,omitempty"`
	Disks           []disk           `yaml:"disks,omitempty"`
}

func optToReqStorage(opt workloadOptions) ([]types.StorageResource, error) {
	storage := make([]types.StorageResource, 0)
	bootableCount := 0
	for _, disk := range opt.Disks {
		res := types.StorageResource{
			Size:      disk.Size,
			Bootable:  disk.Bootable,
			Ephemeral: disk.Ephemeral,
		}

		// Use existing volume
		if disk.ID != nil {
			res.ID = *disk.ID
		} else {
			// Create a new one
			if disk.Source.Type == "" {
				disk.Source.Type = types.Empty
			}

			if disk.Source.Type != types.Empty {
				res.SourceType = disk.Source.Type
				res.SourceID = disk.Source.ID

				if res.SourceID == "" {
					return nil, errors.New("Invalid workload yaml: when using a source an id must also be specified")
				}
			} else {
				if disk.Bootable == true {
					// you may not request a bootable drive
					// from an empty source
					return nil, errors.New("Invalid workload yaml: empty disk source may not be bootable")
				}

				if disk.Size <= 0 {
					return nil, errors.New("Invalid workload yaml: size required when creating a volume")
				}
			}
		}

		if disk.Bootable {
			bootableCount++
		}

		storage = append(storage, res)
	}

	if payloads.Hypervisor(opt.VMType) == payloads.QEMU && bootableCount == 0 {
		return nil, errors.New("Invalid workload yaml: no bootable disks specified for a VM")
	}

	return storage, nil
}

func optToReq(opt workloadOptions, req *types.Workload) error {
	b, err := ioutil.ReadFile(opt.CloudConfigFile)
	if err != nil {
		return err
	}

	config := string(b)

	// this is where you'd validate that the options make
	// sense.
	req.Description = opt.Description
	req.VMType = payloads.Hypervisor(opt.VMType)
	req.FWType = opt.FWType
	req.ImageName = opt.ImageName
	req.Config = config
	req.Storage, err = optToReqStorage(opt)

	if err != nil {
		return err
	}

	// all default resources are required.
	defaults := opt.Defaults

	r := payloads.RequestedResource{
		Type:  payloads.VCPUs,
		Value: defaults.VCPUs,
	}
	req.Defaults = append(req.Defaults, r)

	r = payloads.RequestedResource{
		Type:  payloads.MemMB,
		Value: defaults.MemMB,
	}
	req.Defaults = append(req.Defaults, r)

	return nil
}

func outputWorkload(w types.Workload) {
	var opt workloadOptions

	opt.Description = w.Description
	opt.VMType = string(w.VMType)
	opt.FWType = w.FWType
	opt.ImageName = w.ImageName
	for _, d := range w.Defaults {
		if d.Type == payloads.VCPUs {
			opt.Defaults.VCPUs = d.Value
		} else if d.Type == payloads.MemMB {
			opt.Defaults.MemMB = d.Value
		}
	}

	for _, s := range w.Storage {
		d := disk{
			Size:      s.Size,
			Bootable:  s.Bootable,
			Ephemeral: s.Ephemeral,
		}
		if s.ID != "" {
			d.ID = &s.ID
		}

		src := source{
			Type: s.SourceType,
			ID:   s.SourceID,
		}

		d.Source = src

		opt.Disks = append(opt.Disks, d)
	}

	b, err := yaml.Marshal(opt)
	if err != nil {
		fatalf(err.Error())
	}

	fmt.Println(string(b))
	fmt.Println(w.Config)
}

func (cmd *workloadCreateCommand) run(args []string) error {
	var opt workloadOptions
	var req types.Workload

	if cmd.yamlFile == "" {
		cmd.usage()
	}

	f, err := ioutil.ReadFile(cmd.yamlFile)
	if err != nil {
		fatalf("Unable to read workload config file: %s\n", err)
	}

	err = yaml.Unmarshal(f, &opt)
	if err != nil {
		fatalf("Config file invalid: %s\n", err)
	}

	err = optToReq(opt, &req)
	if err != nil {
		fatalf(err.Error())
	}

	b, err := json.Marshal(req)
	if err != nil {
		fatalf(err.Error())
	}

	body := bytes.NewReader(b)

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	resp, err := sendCiaoRequest("POST", url, nil, body, ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusCreated {
		fatalf("Workload creation failed: %s", resp.Status)
	}

	var workload types.WorkloadResponse

	err = unmarshalHTTPResponse(resp, &workload)
	if err != nil {
		fatalf(err.Error())
	}

	fmt.Printf("Created new workload: %s\n", workload.Workload.ID)

	return nil
}

type workloadDeleteCommand struct {
	Flag     flag.FlagSet
	workload string
}

func (cmd *workloadDeleteCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload delete [flags]

Deletes a given workload

The delete flags are:

`)
	cmd.Flag.PrintDefaults()
	os.Exit(2)
}

func (cmd *workloadDeleteCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.workload, "workload", "", "Workload UUID")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadDeleteCommand) run(args []string) error {
	if cmd.workload == "" {
		cmd.usage()
	}

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	// you should do a get first and search for the workload,
	// then use the href - but not with the currently used
	// OpenStack API. Until we support GET with a ciao API,
	// just hard code the path.
	url = fmt.Sprintf("%s/%s", url, cmd.workload)

	resp, err := sendCiaoRequest("DELETE", url, nil, nil, ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusNoContent {
		fatalf("Workload deletion failed: %s", resp.Status)
	}

	return nil
}

type workloadShowCommand struct {
	Flag     flag.FlagSet
	template string
	workload string
}

func (cmd *workloadShowCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload show

Show workload details

`)
	cmd.Flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n%s",
		tfortools.GenerateUsageDecorated("f", types.Workload{}, nil))
	os.Exit(2)
}

func (cmd *workloadShowCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.workload, "workload", "", "Workload UUID")
	cmd.Flag.StringVar(&cmd.template, "f", "", "Template used to format output")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadShowCommand) run(args []string) error {
	var wl types.Workload

	if cmd.workload == "" {
		cmd.usage()
	}

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	// you should do a get first and search for the workload,
	// then use the href - but not with the currently used
	// OpenStack API. Until we support GET with a ciao API,
	// just hard code the path.
	url = fmt.Sprintf("%s/%s", url, cmd.workload)

	resp, err := sendCiaoRequest("GET", url, nil, nil, ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusOK {
		fatalf("Workload show failed: %s", resp.Status)
	}

	err = unmarshalHTTPResponse(resp, &wl)
	if err != nil {
		fatalf(err.Error())
	}

	if cmd.template != "" {
		return tfortools.OutputToTemplate(os.Stdout, "workload-show", cmd.template, &wl, nil)
	}

	outputWorkload(wl)
	return nil
}
<|endoftext|>


<commit_before>package raphanusclient

import (
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
)

func Test_HTTP(t *testing.T) {
	app := New()

	ts := getTestServer(map[string]string{
		"/test1": `response 1`,
		"/":      `404`,
	})
	defer ts.Close()

	reader, err := app.httpGet(ts.URL + "/test1")
	if err != nil {
		t.Errorf("1. httpGet() failed, error: %s", err)
	}
	body, err := io.ReadAll(reader)
	if err != nil {
		t.Errorf("2. httpGet() failed, error: %s", err)
	}
	if string(body) != "response 1" {
		t.Errorf("3. httpGet() failed, body: %s", string(body))
	}

	reader, err = app.httpPost(ts.URL+"/test1", nil)
	if err != nil {
		t.Errorf("1. httpPost() failed, error: %s", err)
	}
	body, err = io.ReadAll(reader)
	if err != nil {
		t.Errorf("2. httpPost() failed, error: %s", err)
	}
	if string(body) != "response 1" {
		t.Errorf("3. httpPost() failed, body: %s", string(body))
	}

	reader, err = app.httpPut(ts.URL+"/test1", nil)
	if err != nil {
		t.Errorf("1. httpPut() failed, error: %s", err)
	}
	body, err = io.ReadAll(reader)
	if err != nil {
		t.Errorf("2. httpPut() failed, error: %s", err)
	}
	if string(body) != "response 1" {
		t.Errorf("3. httpPut() failed, body: %s", string(body))
	}

	reader, err = app.httpDelete(ts.URL + "/test1")
	if err != nil {
		t.Errorf("1. httpDelete() failed, error: %s", err)
	}
	body, err = io.ReadAll(reader)
	if err != nil {
		t.Errorf("2. httpDelete() failed, error: %s", err)
	}
	if string(body) != "response 1" {
		t.Errorf("3. httpDelete() failed, body: %s", string(body))
	}

	_, err = app.httpClient("FAKE method", "http://example/test3", nil)
	if err == nil {
		t.Errorf("1. callHTTP() failed")
	}
}

func getTestServer(pathMapping map[string]string) *httptest.Server {
	mux := http.NewServeMux()
	for path, result := range pathMapping {
		result := result
		mux.HandleFunc(path, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			_, _ = io.WriteString(w, result)
		}))
	}

	return httptest.NewServer(mux)
}
<commit_msg>Fixed warning<commit_after>package raphanusclient

import (
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
)

func Test_HTTP(t *testing.T) {
	app := New()

	ts := getTestServer(map[string]string{
		"/test1": `response 1`,
		"/":      `404`,
	})
	defer ts.Close()

	reader, err := app.httpGet(ts.URL + "/test1")
	if err != nil {
		t.Errorf("1. httpGet() failed, error: %s", err)
	}
	body, err := io.ReadAll(reader)
	if err != nil {
		t.Errorf("2. httpGet() failed, error: %s", err)
	}
	if string(body) != "response 1" {
		t.Errorf("3. httpGet() failed, body: %s", string(body))
	}

	reader, err = app.httpPost(ts.URL+"/test1", nil)
	if err != nil {
		t.Errorf("1. httpPost() failed, error: %s", err)
	}
	body, err = io.ReadAll(reader)
	if err != nil {
		t.Errorf("2. httpPost() failed, error: %s", err)
	}
	if string(body) != "response 1" {
		t.Errorf("3. httpPost() failed, body: %s", string(body))
	}

	reader, err = app.httpPut(ts.URL+"/test1", nil)
	if err != nil {
		t.Errorf("1. httpPut() failed, error: %s", err)
	}
	body, err = io.ReadAll(reader)
	if err != nil {
		t.Errorf("2. httpPut() failed, error: %s", err)
	}
	if string(body) != "response 1" {
		t.Errorf("3. httpPut() failed, body: %s", string(body))
	}

	reader, err = app.httpDelete(ts.URL + "/test1")
	if err != nil {
		t.Errorf("1. httpDelete() failed, error: %s", err)
	}
	body, err = io.ReadAll(reader)
	if err != nil {
		t.Errorf("2. httpDelete() failed, error: %s", err)
	}
	if string(body) != "response 1" {
		t.Errorf("3. httpDelete() failed, body: %s", string(body))
	}

	_, err = app.httpClient("FAKE method", "http://example/test3", nil)
	if err == nil {
		t.Errorf("1. callHTTP() failed")
	}
}

func getTestServer(pathMapping map[string]string) *httptest.Server {
	mux := http.NewServeMux()
	for path, result := range pathMapping {
		result := result
		mux.HandleFunc(path, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			_, _ = io.WriteString(w, result)
		}))
	}

	return httptest.NewServer(mux)
}
<|endoftext|>


<commit_before>package driver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/helper/args"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/mitchellh/mapstructure"
)

var (
	reRktVersion  = regexp.MustCompile(`rkt version (\d[.\d]+)`)
	reAppcVersion = regexp.MustCompile(`appc version (\d[.\d]+)`)
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
	DriverContext
	fingerprint.StaticFingerprinter
}

type RktDriverConfig struct {
	ImageName string   `mapstructure:"image"`
	Args      []string `mapstructure:"args"`
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
	proc   *os.Process
	image  string
	logger *log.Logger
	waitCh chan *cstructs.WaitResult
	doneCh chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
	Pid   int
	Image string
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
	return &RktDriver{DriverContext: *ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	// Only enable if we are root when running on non-windows systems.
	if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
		d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
		return false, nil
	}

	outBytes, err := exec.Command("rkt", "version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	rktMatches := reRktVersion.FindStringSubmatch(out)
	appcMatches := reAppcVersion.FindStringSubmatch(out)
	if len(rktMatches) != 2 || len(appcMatches) != 2 {
		return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
	}

	node.Attributes["driver.rkt"] = "1"
	node.Attributes["driver.rkt.version"] = rktMatches[1]
	node.Attributes["driver.rkt.appc.version"] = appcMatches[1]

	return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	var driverConfig RktDriverConfig
	if err := mapstructure.WeakDecode(task.Config, &driverConfig); err != nil {
		return nil, err
	}
	// Validate that the config is valid.
	img := driverConfig.ImageName
	if img == "" {
		return nil, fmt.Errorf("Missing ACI image for rkt")
	}

	// Get the tasks local directory.
	taskName := d.DriverContext.taskName
	taskDir, ok := ctx.AllocDir.TaskDirs[taskName]
	if !ok {
		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
	}
	taskLocal := filepath.Join(taskDir, allocdir.TaskLocal)

	// Add the given trust prefix
	trust_prefix, trust_cmd := task.Config["trust_prefix"]
	if trust_cmd {
		var outBuf, errBuf bytes.Buffer
		cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trust_prefix))
		cmd.Stdout = &outBuf
		cmd.Stderr = &errBuf
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf("Error running rkt trust: %s\n\nOutput: %s\n\nError: %s",
				err, outBuf.String(), errBuf.String())
		}
		d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trust_prefix)
	}

	// Build the command.
	var cmd_args []string

	// Inject the environment variables.
	envVars := TaskEnvironmentVariables(ctx, task)

	// Clear the task directories as they are not currently supported.
	envVars.ClearTaskLocalDir()
	envVars.ClearAllocDir()

	for k, v := range envVars.Map() {
		cmd_args = append(cmd_args, fmt.Sprintf("--set-env=%v=%v", k, v))
	}

	// Disble signature verification if the trust command was not run.
	if !trust_cmd {
		cmd_args = append(cmd_args, "--insecure-skip-verify")
	}

	// Append the run command.
	cmd_args = append(cmd_args, "run", "--mds-register=false", img)

	// Check if the user has overriden the exec command.
	if exec_cmd, ok := task.Config["command"]; ok {
		cmd_args = append(cmd_args, fmt.Sprintf("--exec=%v", exec_cmd))
	}

	if task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Memory limit cannot be zero")
	}
	if task.Resources.CPU == 0 {
		return nil, fmt.Errorf("CPU limit cannot be zero")
	}

	// Add memory isolator
	cmd_args = append(cmd_args, fmt.Sprintf("--memory=%vM", int64(task.Resources.MemoryMB)*1024*1024))

	// Add CPU isolator
	cmd_args = append(cmd_args, fmt.Sprintf("--cpu=%vm", int64(task.Resources.CPU)))

	// Add user passed arguments.
	if len(driverConfig.Args) != 0 {
		parsed := args.ParseAndReplace(driverConfig.Args, envVars.Map())

		// Need to start arguments with "--"
		if len(parsed) > 0 {
			cmd_args = append(cmd_args, "--")
		}

		for _, arg := range parsed {
			cmd_args = append(cmd_args, fmt.Sprintf("%v", arg))
		}
	}

	// Create files to capture stdin and out.
	stdoutFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stdout", taskName))
	stderrFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stderr", taskName))

	stdo, err := os.OpenFile(stdoutFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to redirect stdout: %v", err)
	}

	stde, err := os.OpenFile(stderrFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to redirect stderr: %v", err)
	}

	cmd := exec.Command("rkt", cmd_args...)
	cmd.Stdout = stdo
	cmd.Stderr = stde

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("Error running rkt: %v", err)
	}

	d.logger.Printf("[DEBUG] driver.rkt: started ACI %q with: %v", img, cmd.Args)
	h := &rktHandle{
		proc:   cmd.Process,
		image:  img,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan *cstructs.WaitResult, 1),
	}
	go h.run()
	return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	// Parse the handle
	pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
	qpid := &rktPID{}
	if err := json.Unmarshal(pidBytes, qpid); err != nil {
		return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
	}

	// Find the process
	proc, err := os.FindProcess(qpid.Pid)
	if proc == nil || err != nil {
		return nil, fmt.Errorf("failed to find Rkt PID %d: %v", qpid.Pid, err)
	}

	// Return a driver handle
	h := &rktHandle{
		proc:   proc,
		image:  qpid.Image,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan *cstructs.WaitResult, 1),
	}

	go h.run()
	return h, nil
}

func (h *rktHandle) ID() string {
	// Return a handle to the PID
	pid := &rktPID{
		Pid:   h.proc.Pid,
		Image: h.image,
	}
	data, err := json.Marshal(pid)
	if err != nil {
		h.logger.Printf("[ERR] driver.rkt: failed to marshal rkt PID to JSON: %s", err)
	}
	return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan *cstructs.WaitResult {
	return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
	// Update is not possible
	return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
func (h *rktHandle) Kill() error {
	h.proc.Signal(os.Interrupt)
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(5 * time.Second):
		return h.proc.Kill()
	}
}

func (h *rktHandle) run() {
	ps, err := h.proc.Wait()
	close(h.doneCh)
	code := 0
	if !ps.Success() {
		// TODO: Better exit code parsing.
		code = 1
	}
	h.waitCh <- cstructs.NewWaitResult(code, 0, err)
	close(h.waitCh)
}
<commit_msg>Use camelCase for variable names<commit_after>package driver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/helper/args"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/mitchellh/mapstructure"
)

var (
	reRktVersion  = regexp.MustCompile(`rkt version (\d[.\d]+)`)
	reAppcVersion = regexp.MustCompile(`appc version (\d[.\d]+)`)
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
	DriverContext
	fingerprint.StaticFingerprinter
}

type RktDriverConfig struct {
	ImageName string   `mapstructure:"image"`
	Args      []string `mapstructure:"args"`
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
	proc   *os.Process
	image  string
	logger *log.Logger
	waitCh chan *cstructs.WaitResult
	doneCh chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
	Pid   int
	Image string
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
	return &RktDriver{DriverContext: *ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	// Only enable if we are root when running on non-windows systems.
	if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
		d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
		return false, nil
	}

	outBytes, err := exec.Command("rkt", "version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	rktMatches := reRktVersion.FindStringSubmatch(out)
	appcMatches := reAppcVersion.FindStringSubmatch(out)
	if len(rktMatches) != 2 || len(appcMatches) != 2 {
		return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
	}

	node.Attributes["driver.rkt"] = "1"
	node.Attributes["driver.rkt.version"] = rktMatches[1]
	node.Attributes["driver.rkt.appc.version"] = appcMatches[1]

	return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	var driverConfig RktDriverConfig
	if err := mapstructure.WeakDecode(task.Config, &driverConfig); err != nil {
		return nil, err
	}
	// Validate that the config is valid.
	img := driverConfig.ImageName
	if img == "" {
		return nil, fmt.Errorf("Missing ACI image for rkt")
	}

	// Get the tasks local directory.
	taskName := d.DriverContext.taskName
	taskDir, ok := ctx.AllocDir.TaskDirs[taskName]
	if !ok {
		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
	}
	taskLocal := filepath.Join(taskDir, allocdir.TaskLocal)

	// Add the given trust prefix
	trustPrefix, trustCmd := task.Config["trust_prefix"]
	if trustCmd {
		var outBuf, errBuf bytes.Buffer
		cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trustPrefix))
		cmd.Stdout = &outBuf
		cmd.Stderr = &errBuf
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf("Error running rkt trust: %s\n\nOutput: %s\n\nError: %s",
				err, outBuf.String(), errBuf.String())
		}
		d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trustPrefix)
	}

	// Build the command.
	var cmdArgs []string

	// Inject the environment variables.
	envVars := TaskEnvironmentVariables(ctx, task)

	// Clear the task directories as they are not currently supported.
	envVars.ClearTaskLocalDir()
	envVars.ClearAllocDir()

	for k, v := range envVars.Map() {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--set-env=%v=%v", k, v))
	}

	// Disble signature verification if the trust command was not run.
	if !trustCmd {
		cmdArgs = append(cmdArgs, "--insecure-skip-verify")
	}

	// Append the run command.
	cmdArgs = append(cmdArgs, "run", "--mds-register=false", img)

	// Check if the user has overriden the exec command.
	if execCmd, ok := task.Config["command"]; ok {
		cmdArgs = append(cmdArgs, fmt.Sprintf("--exec=%v", execCmd))
	}

	if task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Memory limit cannot be zero")
	}
	if task.Resources.CPU == 0 {
		return nil, fmt.Errorf("CPU limit cannot be zero")
	}

	// Add memory isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--memory=%vM", int64(task.Resources.MemoryMB)*1024*1024))

	// Add CPU isolator
	cmdArgs = append(cmdArgs, fmt.Sprintf("--cpu=%vm", int64(task.Resources.CPU)))

	// Add user passed arguments.
	if len(driverConfig.Args) != 0 {
		parsed := args.ParseAndReplace(driverConfig.Args, envVars.Map())

		// Need to start arguments with "--"
		if len(parsed) > 0 {
			cmdArgs = append(cmdArgs, "--")
		}

		for _, arg := range parsed {
			cmdArgs = append(cmdArgs, fmt.Sprintf("%v", arg))
		}
	}

	// Create files to capture stdin and out.
	stdoutFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stdout", taskName))
	stderrFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stderr", taskName))

	stdo, err := os.OpenFile(stdoutFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to redirect stdout: %v", err)
	}

	stde, err := os.OpenFile(stderrFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to redirect stderr: %v", err)
	}

	cmd := exec.Command("rkt", cmdArgs...)
	cmd.Stdout = stdo
	cmd.Stderr = stde

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("Error running rkt: %v", err)
	}

	d.logger.Printf("[DEBUG] driver.rkt: started ACI %q with: %v", img, cmd.Args)
	h := &rktHandle{
		proc:   cmd.Process,
		image:  img,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan *cstructs.WaitResult, 1),
	}
	go h.run()
	return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	// Parse the handle
	pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
	qpid := &rktPID{}
	if err := json.Unmarshal(pidBytes, qpid); err != nil {
		return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
	}

	// Find the process
	proc, err := os.FindProcess(qpid.Pid)
	if proc == nil || err != nil {
		return nil, fmt.Errorf("failed to find Rkt PID %d: %v", qpid.Pid, err)
	}

	// Return a driver handle
	h := &rktHandle{
		proc:   proc,
		image:  qpid.Image,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan *cstructs.WaitResult, 1),
	}

	go h.run()
	return h, nil
}

func (h *rktHandle) ID() string {
	// Return a handle to the PID
	pid := &rktPID{
		Pid:   h.proc.Pid,
		Image: h.image,
	}
	data, err := json.Marshal(pid)
	if err != nil {
		h.logger.Printf("[ERR] driver.rkt: failed to marshal rkt PID to JSON: %s", err)
	}
	return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan *cstructs.WaitResult {
	return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
	// Update is not possible
	return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
func (h *rktHandle) Kill() error {
	h.proc.Signal(os.Interrupt)
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(5 * time.Second):
		return h.proc.Kill()
	}
}

func (h *rktHandle) run() {
	ps, err := h.proc.Wait()
	close(h.doneCh)
	code := 0
	if !ps.Success() {
		// TODO: Better exit code parsing.
		code = 1
	}
	h.waitCh <- cstructs.NewWaitResult(code, 0, err)
	close(h.waitCh)
}
<|endoftext|>


<commit_before>package stats

import (
	"runtime"
	"time"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/disk"
	"github.com/shirou/gopsutil/host"
	"github.com/shirou/gopsutil/mem"

	shelpers "github.com/hashicorp/nomad/helper/stats"
)

// HostStats represents resource usage stats of the host running a Nomad client
type HostStats struct {
	Memory           *MemoryStats
	CPU              []*CPUStats
	DiskStats        []*DiskStats
	Uptime           uint64
	Timestamp        int64
	CPUTicksConsumed float64
}

// MemoryStats represnts stats related to virtual memory usage
type MemoryStats struct {
	Total     uint64
	Available uint64
	Used      uint64
	Free      uint64
}

// CPUStats represents stats related to cpu usage
type CPUStats struct {
	CPU    string
	User   float64
	System float64
	Idle   float64
	Total  float64
}

// DiskStats represents stats related to disk usage
type DiskStats struct {
	Device            string
	Mountpoint        string
	Size              uint64
	Used              uint64
	Available         uint64
	UsedPercent       float64
	InodesUsedPercent float64
}

// HostStatsCollector collects host resource usage stats
type HostStatsCollector struct {
	clkSpeed        float64
	numCores        int
	statsCalculator map[string]*HostCpuStatsCalculator
}

// NewHostStatsCollector returns a HostStatsCollector
func NewHostStatsCollector() *HostStatsCollector {
	numCores := runtime.NumCPU()
	statsCalculator := make(map[string]*HostCpuStatsCalculator)
	collector := &HostStatsCollector{
		statsCalculator: statsCalculator,
		numCores:        numCores,
	}
	return collector
}

// Collect collects stats related to resource usage of a host
func (h *HostStatsCollector) Collect() (*HostStats, error) {
	hs := &HostStats{Timestamp: time.Now().UTC().UnixNano()}
	if memStats, err := mem.VirtualMemory(); err == nil {
		ms := &MemoryStats{
			Total:     memStats.Total,
			Available: memStats.Available,
			Used:      memStats.Used,
			Free:      memStats.Free,
		}
		hs.Memory = ms
	}

	ticksConsumed := 0.0
	if cpuStats, err := cpu.Times(true); err == nil {
		cs := make([]*CPUStats, len(cpuStats))
		for idx, cpuStat := range cpuStats {
			cs[idx] = &CPUStats{
				CPU:    cpuStat.CPU,
				User:   cpuStat.User,
				System: cpuStat.System,
				Idle:   cpuStat.Idle,
			}
			percentCalculator, ok := h.statsCalculator[cpuStat.CPU]
			if !ok {
				percentCalculator = NewHostCpuStatsCalculator()
				h.statsCalculator[cpuStat.CPU] = percentCalculator
			}
			idle, user, system, total := percentCalculator.Calculate(cpuStat)
			cs[idx].Idle = idle
			cs[idx].System = system
			cs[idx].User = user
			cs[idx].Total = total
			ticksConsumed += (total / 100) * (shelpers.TotalTicksAvailable() / float64(len(cpuStats)))
		}
		hs.CPU = cs
		hs.CPUTicksConsumed = ticksConsumed
	}

	if partitions, err := disk.Partitions(false); err == nil {
		var diskStats []*DiskStats
		for _, partition := range partitions {
			if usage, err := disk.Usage(partition.Mountpoint); err == nil {
				ds := DiskStats{
					Device:            partition.Device,
					Mountpoint:        partition.Mountpoint,
					Size:              usage.Total,
					Used:              usage.Used,
					Available:         usage.Free,
					UsedPercent:       usage.UsedPercent,
					InodesUsedPercent: usage.InodesUsedPercent,
				}
				diskStats = append(diskStats, &ds)
			}
		}
		hs.DiskStats = diskStats
	}

	if uptime, err := host.Uptime(); err == nil {
		hs.Uptime = uptime
	}
	return hs, nil
}

// HostCpuStatsCalculator calculates cpu usage percentages
type HostCpuStatsCalculator struct {
	prevIdle   float64
	prevUser   float64
	prevSystem float64
	prevBusy   float64
	prevTotal  float64
}

// NewHostCpuStatsCalculator returns a HostCpuStatsCalculator
func NewHostCpuStatsCalculator() *HostCpuStatsCalculator {
	return &HostCpuStatsCalculator{}
}

// Calculate calculates the current cpu usage percentages
func (h *HostCpuStatsCalculator) Calculate(times cpu.TimesStat) (idle float64, user float64, system float64, total float64) {
	currentIdle := times.Idle
	currentUser := times.User
	currentSystem := times.System
	currentTotal := times.Total()

	deltaTotal := currentTotal - h.prevTotal
	idle = ((currentIdle - h.prevIdle) / deltaTotal) * 100
	user = ((currentUser - h.prevUser) / deltaTotal) * 100
	system = ((currentSystem - h.prevSystem) / deltaTotal) * 100

	currentBusy := times.User + times.System + times.Nice + times.Iowait + times.Irq +
		times.Softirq + times.Steal + times.Guest + times.GuestNice + times.Stolen

	total = ((currentBusy - h.prevBusy) / deltaTotal) * 100

	h.prevIdle = currentIdle
	h.prevUser = currentUser
	h.prevSystem = currentSystem
	h.prevTotal = currentTotal
	h.prevBusy = currentBusy

	return
}
<commit_msg>Removing un-used code<commit_after>package stats

import (
	"runtime"
	"time"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/disk"
	"github.com/shirou/gopsutil/host"
	"github.com/shirou/gopsutil/mem"

	shelpers "github.com/hashicorp/nomad/helper/stats"
)

// HostStats represents resource usage stats of the host running a Nomad client
type HostStats struct {
	Memory           *MemoryStats
	CPU              []*CPUStats
	DiskStats        []*DiskStats
	Uptime           uint64
	Timestamp        int64
	CPUTicksConsumed float64
}

// MemoryStats represnts stats related to virtual memory usage
type MemoryStats struct {
	Total     uint64
	Available uint64
	Used      uint64
	Free      uint64
}

// CPUStats represents stats related to cpu usage
type CPUStats struct {
	CPU    string
	User   float64
	System float64
	Idle   float64
	Total  float64
}

// DiskStats represents stats related to disk usage
type DiskStats struct {
	Device            string
	Mountpoint        string
	Size              uint64
	Used              uint64
	Available         uint64
	UsedPercent       float64
	InodesUsedPercent float64
}

// HostStatsCollector collects host resource usage stats
type HostStatsCollector struct {
	clkSpeed        float64
	numCores        int
	statsCalculator map[string]*HostCpuStatsCalculator
}

// NewHostStatsCollector returns a HostStatsCollector
func NewHostStatsCollector() *HostStatsCollector {
	numCores := runtime.NumCPU()
	statsCalculator := make(map[string]*HostCpuStatsCalculator)
	collector := &HostStatsCollector{
		statsCalculator: statsCalculator,
		numCores:        numCores,
	}
	return collector
}

// Collect collects stats related to resource usage of a host
func (h *HostStatsCollector) Collect() (*HostStats, error) {
	hs := &HostStats{Timestamp: time.Now().UTC().UnixNano()}
	if memStats, err := mem.VirtualMemory(); err == nil {
		ms := &MemoryStats{
			Total:     memStats.Total,
			Available: memStats.Available,
			Used:      memStats.Used,
			Free:      memStats.Free,
		}
		hs.Memory = ms
	}

	ticksConsumed := 0.0
	if cpuStats, err := cpu.Times(true); err == nil {
		cs := make([]*CPUStats, len(cpuStats))
		for idx, cpuStat := range cpuStats {
			percentCalculator, ok := h.statsCalculator[cpuStat.CPU]
			if !ok {
				percentCalculator = NewHostCpuStatsCalculator()
				h.statsCalculator[cpuStat.CPU] = percentCalculator
			}
			idle, user, system, total := percentCalculator.Calculate(cpuStat)
			cs[idx] = &CPUStats{
				CPU:    cpuStat.CPU,
				User:   user,
				System: system,
				Idle:   idle,
				Total:  total,
			}
			ticksConsumed += (total / 100) * (shelpers.TotalTicksAvailable() / float64(len(cpuStats)))
		}
		hs.CPU = cs
		hs.CPUTicksConsumed = ticksConsumed
	}

	if partitions, err := disk.Partitions(false); err == nil {
		var diskStats []*DiskStats
		for _, partition := range partitions {
			if usage, err := disk.Usage(partition.Mountpoint); err == nil {
				ds := DiskStats{
					Device:            partition.Device,
					Mountpoint:        partition.Mountpoint,
					Size:              usage.Total,
					Used:              usage.Used,
					Available:         usage.Free,
					UsedPercent:       usage.UsedPercent,
					InodesUsedPercent: usage.InodesUsedPercent,
				}
				diskStats = append(diskStats, &ds)
			}
		}
		hs.DiskStats = diskStats
	}

	if uptime, err := host.Uptime(); err == nil {
		hs.Uptime = uptime
	}
	return hs, nil
}

// HostCpuStatsCalculator calculates cpu usage percentages
type HostCpuStatsCalculator struct {
	prevIdle   float64
	prevUser   float64
	prevSystem float64
	prevBusy   float64
	prevTotal  float64
}

// NewHostCpuStatsCalculator returns a HostCpuStatsCalculator
func NewHostCpuStatsCalculator() *HostCpuStatsCalculator {
	return &HostCpuStatsCalculator{}
}

// Calculate calculates the current cpu usage percentages
func (h *HostCpuStatsCalculator) Calculate(times cpu.TimesStat) (idle float64, user float64, system float64, total float64) {
	currentIdle := times.Idle
	currentUser := times.User
	currentSystem := times.System
	currentTotal := times.Total()

	deltaTotal := currentTotal - h.prevTotal
	idle = ((currentIdle - h.prevIdle) / deltaTotal) * 100
	user = ((currentUser - h.prevUser) / deltaTotal) * 100
	system = ((currentSystem - h.prevSystem) / deltaTotal) * 100

	currentBusy := times.User + times.System + times.Nice + times.Iowait + times.Irq +
		times.Softirq + times.Steal + times.Guest + times.GuestNice + times.Stolen

	total = ((currentBusy - h.prevBusy) / deltaTotal) * 100

	h.prevIdle = currentIdle
	h.prevUser = currentUser
	h.prevSystem = currentSystem
	h.prevTotal = currentTotal
	h.prevBusy = currentBusy

	return
}
<|endoftext|>


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

// Package https provides a helper for starting an HTTPS server.
package https

import (
	"context"
	"crypto/tls"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"

	"google.golang.org/api/option"

	"cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/storage"
	"rsc.io/letsencrypt"

	"upspin.io/auth"
	"upspin.io/log"
)

// Options permits the configuration of TLS certificates for servers running
// outside GCE. The default is the self-signed certificate in
// upspin.io/auth/grpcauth/testdata.
type Options struct {
	CertFile string
	KeyFile  string
}

var defaultOptions = &Options{
	CertFile: filepath.Join(os.Getenv("GOPATH"), "/src/upspin.io/auth/grpcauth/testdata/cert.pem"),
	KeyFile:  filepath.Join(os.Getenv("GOPATH"), "/src/upspin.io/auth/grpcauth/testdata/key.pem"),
}

func (opt *Options) applyDefaults() {
	if opt.CertFile == "" {
		opt.CertFile = defaultOptions.CertFile
	}
	if opt.KeyFile == "" {
		opt.KeyFile = defaultOptions.KeyFile
	}
}

// ListenAndServe serves the http.DefaultServeMux by HTTPS (and HTTP,
// redirecting to HTTPS), storing SSL credentials in the Google Cloud Storage
// buckets nominated by the Google Compute Engine project metadata variables
// "letscloud-get-url-metaSuffix" and "letscloud-put-url-metaSuffix", where
// metaSuffix is the supplied argument.
// (See the upspin.io/cloud/letscloud package for more information.)
//
// If the server is running outside GCE, instead an HTTPS server is started on
// the address specified by addr the certificate and key files specified by
// opt.
func ListenAndServe(metaSuffix, addr string, opt *Options) {
	if opt == nil {
		opt = defaultOptions
	} else {
		opt.applyDefaults()
	}
	if metadata.OnGCE() {
		log.Info.Println("https: on GCE; serving HTTPS on port 443 using Let's Encrypt")
		var m letsencrypt.Manager
		const key = "letsencrypt-bucket"
		bucket, err := metadata.InstanceAttributeValue(key)
		if err != nil {
			log.Fatalf("https: couldn't read %q metadata value: %v", key, err)
		}
		if err := letsencryptCache(&m, bucket, metaSuffix); err != nil {
			log.Fatalf("https: couldn't set up letsencrypt cache: %v", err)
		}
		log.Fatalf("https: %v", m.Serve())
	}

	log.Info.Printf("https: not on GCE; serving HTTPS on %q", addr)
	if opt.CertFile == defaultOptions.CertFile || opt.KeyFile == defaultOptions.KeyFile {
		log.Error.Print("https: WARNING: using self-signed test certificates.")
	}
	config, err := auth.NewDefaultTLSConfig(opt.CertFile, opt.KeyFile)
	if err != nil {
		log.Fatalf("https: setting up TLS config: %v", err)
	}
	config.NextProtos = []string{"h2"} // Enable HTTP/2 support
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatalf("https: %v", err)
	}
	err = http.Serve(tls.NewListener(ln, config), nil)
	log.Fatalf("https: %v", err)
}

func letsencryptCache(m *letsencrypt.Manager, bucket, suffix string) error {
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithScopes(storage.ScopeFullControl))
	if err != nil {
		return err
	}
	obj := client.Bucket(bucket).Object("letsencrypt-" + suffix)

	// Try to read the existing cache value, if present.
	r, err := obj.NewReader(ctx)
	if err != storage.ErrObjectNotExist {
		if err != nil {
			return err
		}
		data, err := ioutil.ReadAll(r)
		r.Close()
		if err != nil {
			return err
		}
		if err := m.Unmarshal(string(data)); err != nil {
			return err
		}
	}

	go func() {
		// Watch the letsencrypt manager for changes and cache them.
		for range m.Watch() {
			w := obj.NewWriter(ctx)
			_, err := io.WriteString(w, m.Marshal())
			if err != nil {
				log.Printf("https: writing letsencrypt cache: %v", err)
				continue
			}
			if err := w.Close(); err != nil {
				log.Printf("https: writing letsencrypt cache: %v", err)
			}
		}
	}()

	return nil
}
<commit_msg>cloud/https: add option to run Let's Encrypt when not on GCE<commit_after>// Copyright 2016 The Upspin Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package https provides a helper for starting an HTTPS server.
package https

import (
	"context"
	"crypto/tls"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"

	"google.golang.org/api/option"

	"cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/storage"
	"rsc.io/letsencrypt"

	"upspin.io/auth"
	"upspin.io/log"
)

// Options permits the configuration of TLS certificates for servers running
// outside GCE. The default is the self-signed certificate in
// upspin.io/auth/grpcauth/testdata.
type Options struct {
	// LetsEncryptCache specifies the cache file for Let's Encrypt.
	// If set, enables Let's Encrypt certificates for this server.
	LetsEncryptCache string

	// CertFile and KeyFile specifies the TLS certificates to use.
	// It has no effect if LetsEncryptCache is set.
	CertFile string
	KeyFile  string
}

var defaultOptions = &Options{
	CertFile: filepath.Join(os.Getenv("GOPATH"), "/src/upspin.io/auth/grpcauth/testdata/cert.pem"),
	KeyFile:  filepath.Join(os.Getenv("GOPATH"), "/src/upspin.io/auth/grpcauth/testdata/key.pem"),
}

func (opt *Options) applyDefaults() {
	if opt.CertFile == "" {
		opt.CertFile = defaultOptions.CertFile
	}
	if opt.KeyFile == "" {
		opt.KeyFile = defaultOptions.KeyFile
	}
}

// ListenAndServe serves the http.DefaultServeMux by HTTPS (and HTTP,
// redirecting to HTTPS), storing SSL credentials in the Google Cloud Storage
// buckets nominated by the Google Compute Engine project metadata variables
// "letscloud-get-url-metaSuffix" and "letscloud-put-url-metaSuffix", where
// metaSuffix is the supplied argument.
// (See the upspin.io/cloud/letscloud package for more information.)
//
// If the server is running outside GCE, instead an HTTPS server is started on
// the address specified by addr using the certificate details specified by opt.
func ListenAndServe(metaSuffix, addr string, opt *Options) {
	if opt == nil {
		opt = defaultOptions
	} else {
		opt.applyDefaults()
	}
	if metadata.OnGCE() {
		log.Info.Println("https: on GCE; serving HTTPS on port 443 using Let's Encrypt")
		var m letsencrypt.Manager
		const key = "letsencrypt-bucket"
		bucket, err := metadata.InstanceAttributeValue(key)
		if err != nil {
			log.Fatalf("https: couldn't read %q metadata value: %v", key, err)
		}
		if err := letsencryptCache(&m, bucket, metaSuffix); err != nil {
			log.Fatalf("https: couldn't set up letsencrypt cache: %v", err)
		}
		log.Fatalf("https: %v", m.Serve())
	}

	var config *tls.Config
	if file := opt.LetsEncryptCache; file != "" {
		log.Info.Printf("https: serving HTTPS on %q using Let's Encrypt certificates", addr)
		var m letsencrypt.Manager
		if err := m.CacheFile(file); err != nil {
			log.Fatalf("https: couldn't set up letsencrypt cache: %v", err)
		}
		config = &tls.Config{
			GetCertificate: m.GetCertificate,
		}
	} else {
		log.Info.Printf("https: not on GCE; serving HTTPS on %q using provided certificates", addr)
		if opt.CertFile == defaultOptions.CertFile || opt.KeyFile == defaultOptions.KeyFile {
			log.Error.Print("https: WARNING: using self-signed test certificates.")
		}
		var err error
		config, err = auth.NewDefaultTLSConfig(opt.CertFile, opt.KeyFile)
		if err != nil {
			log.Fatalf("https: setting up TLS config: %v", err)
		}
	}
	config.NextProtos = []string{"h2"} // Enable HTTP/2 support
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatalf("https: %v", err)
	}
	err = http.Serve(tls.NewListener(ln, config), nil)
	log.Fatalf("https: %v", err)
}

func letsencryptCache(m *letsencrypt.Manager, bucket, suffix string) error {
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithScopes(storage.ScopeFullControl))
	if err != nil {
		return err
	}
	obj := client.Bucket(bucket).Object("letsencrypt-" + suffix)

	// Try to read the existing cache value, if present.
	r, err := obj.NewReader(ctx)
	if err != storage.ErrObjectNotExist {
		if err != nil {
			return err
		}
		data, err := ioutil.ReadAll(r)
		r.Close()
		if err != nil {
			return err
		}
		if err := m.Unmarshal(string(data)); err != nil {
			return err
		}
	}

	go func() {
		// Watch the letsencrypt manager for changes and cache them.
		for range m.Watch() {
			w := obj.NewWriter(ctx)
			_, err := io.WriteString(w, m.Marshal())
			if err != nil {
				log.Printf("https: writing letsencrypt cache: %v", err)
				continue
			}
			if err := w.Close(); err != nil {
				log.Printf("https: writing letsencrypt cache: %v", err)
			}
		}
	}()

	return nil
}
<|endoftext|>