<commit_before>package tsa_test

import (
	"context"
	"encoding/json"
	"net/http"
	"time"

	"code.cloudfoundry.org/clock/fakeclock"
	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/garden/gardenfakes"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagerctx"
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/concourse/baggageclaim"
	"github.com/concourse/baggageclaim/baggageclaimfakes"
	"github.com/concourse/concourse/atc"
	. "github.com/concourse/concourse/tsa"
	"github.com/concourse/concourse/tsa/tsafakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/ghttp"
	"github.com/tedsuo/rata"
)

var _ = Describe("Heartbeater", func() {
	type registration struct {
		worker atc.Worker
		ttl    time.Duration
	}

	var (
		ctx    context.Context
		cancel func()

		addrToRegister string
		fakeClock      *fakeclock.FakeClock
		interval       time.Duration
		cprInterval    time.Duration
		resourceTypes  []atc.WorkerResourceType

		expectedWorker         atc.Worker
		fakeTokenGenerator     *tsafakes.FakeTokenGenerator
		fakeGardenClient       *gardenfakes.FakeClient
		fakeBaggageclaimClient *baggageclaimfakes.FakeClient
		fakeATC1               *ghttp.Server
		fakeATC2               *ghttp.Server
		atcEndpointPicker      *tsafakes.FakeEndpointPicker
		heartbeatErr           chan error

		verifyRegister  http.HandlerFunc
		verifyHeartbeat http.HandlerFunc

		registrations <-chan registration
		heartbeats    <-chan registration
		clientWriter  *gbytes.Buffer

		worker atc.Worker
	)

	BeforeEach(func() {
		ctx, cancel = context.WithCancel(lagerctx.NewContext(context.Background(), lagertest.NewTestLogger("test")))

		addrToRegister = "1.2.3.4:7777"
		fakeClock = fakeclock.NewFakeClock(time.Unix(123, 456))
		interval = time.Second
		cprInterval = 100 * time.Millisecond
		resourceTypes = []atc.WorkerResourceType{
			{
				Type:  "git",
				Image: "/path/to/git/resource",
			},
		}

		worker = atc.Worker{
			Name:          "some-name",
			GardenAddr:    addrToRegister,
			ResourceTypes: resourceTypes,
			Platform:      "some-platform",
			Tags:          []string{"some", "tags"},
		}

		expectedWorker = worker

		fakeATC1 = ghttp.NewServer()
		fakeATC2 = ghttp.NewServer()

		registerRoute, found := atc.Routes.FindRouteByName(atc.RegisterWorker)
		Expect(found).To(BeTrue())

		registered := make(chan registration, 100)
		registrations = registered

		heartbeated := make(chan registration, 100)
		heartbeats = heartbeated

		verifyRegister = ghttp.CombineHandlers(
			ghttp.VerifyRequest(registerRoute.Method, registerRoute.Path),
			func(w http.ResponseWriter, r *http.Request) {
				var worker atc.Worker
				Expect(r.Header.Get("Authorization")).To(Equal("Bearer yo"))

				err := json.NewDecoder(r.Body).Decode(&worker)
				Expect(err).NotTo(HaveOccurred())

				ttl, err := time.ParseDuration(r.URL.Query().Get("ttl"))
				Expect(err).NotTo(HaveOccurred())

				registered <- registration{worker, ttl}
			},
		)

		verifyHeartbeat = ghttp.CombineHandlers(
			ghttp.VerifyRequest("PUT", "/api/v1/workers/some-name/heartbeat"),
			func(w http.ResponseWriter, r *http.Request) {
				var worker atc.Worker
				Expect(r.Header.Get("Authorization")).To(Equal("Bearer yo"))

				err := json.NewDecoder(r.Body).Decode(&worker)
				Expect(err).NotTo(HaveOccurred())

				ttl, err := time.ParseDuration(r.URL.Query().Get("ttl"))
				Expect(err).NotTo(HaveOccurred())

				heartbeated <- registration{worker, ttl}

				json.NewEncoder(w).Encode(worker)
			},
		)

		fakeGardenClient = new(gardenfakes.FakeClient)
		fakeBaggageclaimClient = new(baggageclaimfakes.FakeClient)
		fakeTokenGenerator = new(tsafakes.FakeTokenGenerator)

		fakeTokenGenerator.GenerateSystemTokenReturns("yo", nil)
		fakeTokenGenerator.GenerateTeamTokenReturns("yo", nil)
		clientWriter = gbytes.NewBuffer()

		pickCallCount := 0
		atcEndpointPicker = new(tsafakes.FakeEndpointPicker)
		atcEndpointPicker.PickStub = func() *rata.RequestGenerator {
			pickCallCount++

			if pickCallCount%2 == 0 {
				return rata.NewRequestGenerator(fakeATC2.URL(), atc.Routes)
			}

			return rata.NewRequestGenerator(fakeATC1.URL(), atc.Routes)
		}

		heartbeatErr = make(chan error, 1)
	})

	JustBeforeEach(func() {
		heartbeater := NewHeartbeater(
			fakeClock,
			interval,
			cprInterval,
			fakeGardenClient,
			fakeBaggageclaimClient,
			atcEndpointPicker,
			fakeTokenGenerator,
			worker,
			NewEventWriter(clientWriter),
		)

		go func() {
			heartbeatErr <- heartbeater.Heartbeat(ctx)
			close(heartbeatErr)
		}()
	})

	AfterEach(func() {
		cancel()
		<-heartbeatErr
		fakeATC2.Close()
		fakeATC1.Close()
	})

	Context("when Garden returns containers and Baggageclaim returns volumes", func() {
		BeforeEach(func() {
			containers := make(chan []garden.Container, 4)
			volumes := make(chan []baggageclaim.Volume, 4)

			containers <- []garden.Container{
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
			}

			volumes <- []baggageclaim.Volume{
				new(baggageclaimfakes.FakeVolume),
				new(baggageclaimfakes.FakeVolume),
				new(baggageclaimfakes.FakeVolume),
			}

			containers <- []garden.Container{
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
			}

			volumes <- []baggageclaim.Volume{
				new(baggageclaimfakes.FakeVolume),
				new(baggageclaimfakes.FakeVolume),
			}

			containers <- []garden.Container{
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
			}

			volumes <- []baggageclaim.Volume{
				new(baggageclaimfakes.FakeVolume),
			}

			containers <- []garden.Container{
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
			}

			volumes <- []baggageclaim.Volume{}

			close(containers)
			close(volumes)

			fakeGardenClient.ContainersStub = func(garden.Properties) ([]garden.Container, error) {
				return <-containers, nil
			}

			fakeBaggageclaimClient.ListVolumesStub = func(lager.Logger, baggageclaim.VolumeProperties) (baggageclaim.Volumes, error) {
				return <-volumes, nil
			}
		})

		Context("when the ATC responds to registration requests", func() {
			Context("When the DEBUG log level is set", func() {
				BeforeEach(func() {
					fakeATC1.AppendHandlers(verifyRegister)
					fakeATC2.AppendHandlers(verifyHeartbeat)
				})

				It("immediately registers", func() {
					expectedWorker.ActiveContainers = 2
					expectedWorker.ActiveVolumes = 3
					Eventually(registrations).Should(Receive(Equal(registration{expectedWorker, 2 * interval})))
				})

				It("heartbeats", func() {
					Eventually(registrations).Should(Receive())

					fakeClock.WaitForWatcherAndIncrement(interval)
					expectedWorker.ActiveContainers = 5
					expectedWorker.ActiveVolumes = 2
					Eventually(heartbeats).Should(Receive(Equal(registration{expectedWorker, 2 * interval})))
				})

				It("emits events", func() {
					Eventually(registrations).Should(Receive())

					Eventually(clientWriter).Should(gbytes.Say(`{"event":"registered"}`))

					fakeClock.WaitForWatcherAndIncrement(interval)
					Eventually(clientWriter).Should(gbytes.Say(`{"event":"heartbeated"}`))
				})
			})
		})

		Context("when heartbeat returns worker is landed", func() {
			BeforeEach(func() {
				heartbeated := make(chan registration, 100)
				heartbeats = heartbeated

				fakeATC1.AppendHandlers(verifyRegister)
				fakeATC2.AppendHandlers(ghttp.CombineHandlers(
					ghttp.VerifyRequest("PUT", "/api/v1/workers/some-name/heartbeat"),
					func(w http.ResponseWriter, r *http.Request) {
						var worker atc.Worker
						Expect(r.Header.Get("Authorization")).To(Equal("Bearer yo"))

						err := json.NewDecoder(r.Body).Decode(&worker)
						Expect(err).NotTo(HaveOccurred())

						ttl, err := time.ParseDuration(r.URL.Query().Get("ttl"))
						Expect(err).NotTo(HaveOccurred())

						heartbeated <- registration{worker, ttl}

						json.NewEncoder(w).Encode(atc.Worker{
							State: "landed",
						})
					},
				))
			})

			It("exits heartbeater with no error", func() {
				Eventually(registrations).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(interval)
				Eventually(heartbeats).Should(Receive())

				err := <-heartbeatErr
				Expect(err).To(BeNil())
			})
		})

		Context("when the ATC doesn't respond to the first heartbeat", func() {
			BeforeEach(func() {
				fakeATC1.AppendHandlers(
					verifyRegister,
					verifyHeartbeat,
				)
				fakeATC2.AppendHandlers(
					ghttp.CombineHandlers(
						verifyHeartbeat,
						func(w http.ResponseWriter, r *http.Request) { fakeATC2.CloseClientConnections() },
					),
					verifyHeartbeat,
				)
			})

			It("heartbeats faster according to cprInterval", func() {
				Eventually(registrations).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(interval)
				Eventually(heartbeats).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(cprInterval)
				expectedWorker.ActiveContainers = 4
				expectedWorker.ActiveVolumes = 1
				Eventually(heartbeats).Should(Receive(Equal(registration{expectedWorker, 2 * interval})))
			})

			It("goes back to normal after the heartbeat succeeds", func() {
				Eventually(registrations).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(interval)
				Eventually(heartbeats).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(cprInterval)
				Eventually(heartbeats).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(cprInterval)
				Consistently(heartbeats).ShouldNot(Receive())

				fakeClock.WaitForWatcherAndIncrement(interval - cprInterval)
				expectedWorker.ActiveContainers = 3
				expectedWorker.ActiveVolumes = 0
				Eventually(heartbeats).Should(Receive(Equal(registration{expectedWorker, 2 * interval})))
			})
		})
	})
})
<commit_msg>tsa: fix test data race<commit_after>package tsa_test

import (
	"context"
	"encoding/json"
	"net/http"
	"time"

	"code.cloudfoundry.org/clock/fakeclock"
	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/garden/gardenfakes"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagerctx"
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/concourse/baggageclaim"
	"github.com/concourse/baggageclaim/baggageclaimfakes"
	"github.com/concourse/concourse/atc"
	. "github.com/concourse/concourse/tsa"
	"github.com/concourse/concourse/tsa/tsafakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/ghttp"
	"github.com/tedsuo/rata"
)

var _ = Describe("Heartbeater", func() {
	type registration struct {
		worker atc.Worker
		ttl    time.Duration
	}

	var (
		ctx    context.Context
		cancel func()

		addrToRegister string
		fakeClock      *fakeclock.FakeClock
		interval       time.Duration
		cprInterval    time.Duration
		resourceTypes  []atc.WorkerResourceType

		expectedWorker         atc.Worker
		fakeTokenGenerator     *tsafakes.FakeTokenGenerator
		fakeGardenClient       *gardenfakes.FakeClient
		fakeBaggageclaimClient *baggageclaimfakes.FakeClient
		fakeATC1               *ghttp.Server
		fakeATC2               *ghttp.Server
		atcEndpointPicker      *tsafakes.FakeEndpointPicker
		heartbeatErr           <-chan error

		verifyRegister  http.HandlerFunc
		verifyHeartbeat http.HandlerFunc

		registrations <-chan registration
		heartbeats    <-chan registration
		clientWriter  *gbytes.Buffer

		worker atc.Worker
	)

	BeforeEach(func() {
		ctx, cancel = context.WithCancel(lagerctx.NewContext(context.Background(), lagertest.NewTestLogger("test")))

		addrToRegister = "1.2.3.4:7777"
		fakeClock = fakeclock.NewFakeClock(time.Unix(123, 456))
		interval = time.Second
		cprInterval = 100 * time.Millisecond
		resourceTypes = []atc.WorkerResourceType{
			{
				Type:  "git",
				Image: "/path/to/git/resource",
			},
		}

		worker = atc.Worker{
			Name:          "some-name",
			GardenAddr:    addrToRegister,
			ResourceTypes: resourceTypes,
			Platform:      "some-platform",
			Tags:          []string{"some", "tags"},
		}

		expectedWorker = worker

		fakeATC1 = ghttp.NewServer()
		fakeATC2 = ghttp.NewServer()

		registerRoute, found := atc.Routes.FindRouteByName(atc.RegisterWorker)
		Expect(found).To(BeTrue())

		registered := make(chan registration, 100)
		registrations = registered

		heartbeated := make(chan registration, 100)
		heartbeats = heartbeated

		verifyRegister = ghttp.CombineHandlers(
			ghttp.VerifyRequest(registerRoute.Method, registerRoute.Path),
			func(w http.ResponseWriter, r *http.Request) {
				var worker atc.Worker
				Expect(r.Header.Get("Authorization")).To(Equal("Bearer yo"))

				err := json.NewDecoder(r.Body).Decode(&worker)
				Expect(err).NotTo(HaveOccurred())

				ttl, err := time.ParseDuration(r.URL.Query().Get("ttl"))
				Expect(err).NotTo(HaveOccurred())

				registered <- registration{worker, ttl}
			},
		)

		verifyHeartbeat = ghttp.CombineHandlers(
			ghttp.VerifyRequest("PUT", "/api/v1/workers/some-name/heartbeat"),
			func(w http.ResponseWriter, r *http.Request) {
				var worker atc.Worker
				Expect(r.Header.Get("Authorization")).To(Equal("Bearer yo"))

				err := json.NewDecoder(r.Body).Decode(&worker)
				Expect(err).NotTo(HaveOccurred())

				ttl, err := time.ParseDuration(r.URL.Query().Get("ttl"))
				Expect(err).NotTo(HaveOccurred())

				heartbeated <- registration{worker, ttl}

				json.NewEncoder(w).Encode(worker)
			},
		)

		fakeGardenClient = new(gardenfakes.FakeClient)
		fakeBaggageclaimClient = new(baggageclaimfakes.FakeClient)
		fakeTokenGenerator = new(tsafakes.FakeTokenGenerator)

		fakeTokenGenerator.GenerateSystemTokenReturns("yo", nil)
		fakeTokenGenerator.GenerateTeamTokenReturns("yo", nil)
		clientWriter = gbytes.NewBuffer()

		pickCallCount := 0
		atcEndpointPicker = new(tsafakes.FakeEndpointPicker)
		atcEndpointPicker.PickStub = func() *rata.RequestGenerator {
			pickCallCount++

			if pickCallCount%2 == 0 {
				return rata.NewRequestGenerator(fakeATC2.URL(), atc.Routes)
			}

			return rata.NewRequestGenerator(fakeATC1.URL(), atc.Routes)
		}

	})

	JustBeforeEach(func() {
		heartbeater := NewHeartbeater(
			fakeClock,
			interval,
			cprInterval,
			fakeGardenClient,
			fakeBaggageclaimClient,
			atcEndpointPicker,
			fakeTokenGenerator,
			worker,
			NewEventWriter(clientWriter),
		)

		errs := make(chan error, 1)
		heartbeatErr = errs
		go func() {
			errs <- heartbeater.Heartbeat(ctx)
			close(errs)
		}()
	})

	AfterEach(func() {
		cancel()
		<-heartbeatErr
		fakeATC2.Close()
		fakeATC1.Close()
	})

	Context("when Garden returns containers and Baggageclaim returns volumes", func() {
		BeforeEach(func() {
			containers := make(chan []garden.Container, 4)
			volumes := make(chan []baggageclaim.Volume, 4)

			containers <- []garden.Container{
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
			}

			volumes <- []baggageclaim.Volume{
				new(baggageclaimfakes.FakeVolume),
				new(baggageclaimfakes.FakeVolume),
				new(baggageclaimfakes.FakeVolume),
			}

			containers <- []garden.Container{
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
			}

			volumes <- []baggageclaim.Volume{
				new(baggageclaimfakes.FakeVolume),
				new(baggageclaimfakes.FakeVolume),
			}

			containers <- []garden.Container{
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
			}

			volumes <- []baggageclaim.Volume{
				new(baggageclaimfakes.FakeVolume),
			}

			containers <- []garden.Container{
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
				new(gardenfakes.FakeContainer),
			}

			volumes <- []baggageclaim.Volume{}

			close(containers)
			close(volumes)

			fakeGardenClient.ContainersStub = func(garden.Properties) ([]garden.Container, error) {
				return <-containers, nil
			}

			fakeBaggageclaimClient.ListVolumesStub = func(lager.Logger, baggageclaim.VolumeProperties) (baggageclaim.Volumes, error) {
				return <-volumes, nil
			}
		})

		Context("when the ATC responds to registration requests", func() {
			Context("When the DEBUG log level is set", func() {
				BeforeEach(func() {
					fakeATC1.AppendHandlers(verifyRegister)
					fakeATC2.AppendHandlers(verifyHeartbeat)
				})

				It("immediately registers", func() {
					expectedWorker.ActiveContainers = 2
					expectedWorker.ActiveVolumes = 3
					Eventually(registrations).Should(Receive(Equal(registration{expectedWorker, 2 * interval})))
				})

				It("heartbeats", func() {
					Eventually(registrations).Should(Receive())

					fakeClock.WaitForWatcherAndIncrement(interval)
					expectedWorker.ActiveContainers = 5
					expectedWorker.ActiveVolumes = 2
					Eventually(heartbeats).Should(Receive(Equal(registration{expectedWorker, 2 * interval})))
				})

				It("emits events", func() {
					Eventually(registrations).Should(Receive())

					Eventually(clientWriter).Should(gbytes.Say(`{"event":"registered"}`))

					fakeClock.WaitForWatcherAndIncrement(interval)
					Eventually(clientWriter).Should(gbytes.Say(`{"event":"heartbeated"}`))
				})
			})
		})

		Context("when heartbeat returns worker is landed", func() {
			BeforeEach(func() {
				heartbeated := make(chan registration, 100)
				heartbeats = heartbeated

				fakeATC1.AppendHandlers(verifyRegister)
				fakeATC2.AppendHandlers(ghttp.CombineHandlers(
					ghttp.VerifyRequest("PUT", "/api/v1/workers/some-name/heartbeat"),
					func(w http.ResponseWriter, r *http.Request) {
						var worker atc.Worker
						Expect(r.Header.Get("Authorization")).To(Equal("Bearer yo"))

						err := json.NewDecoder(r.Body).Decode(&worker)
						Expect(err).NotTo(HaveOccurred())

						ttl, err := time.ParseDuration(r.URL.Query().Get("ttl"))
						Expect(err).NotTo(HaveOccurred())

						heartbeated <- registration{worker, ttl}

						json.NewEncoder(w).Encode(atc.Worker{
							State: "landed",
						})
					},
				))
			})

			It("exits heartbeater with no error", func() {
				Eventually(registrations).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(interval)
				Eventually(heartbeats).Should(Receive())

				err := <-heartbeatErr
				Expect(err).To(BeNil())
			})
		})

		Context("when the ATC doesn't respond to the first heartbeat", func() {
			BeforeEach(func() {
				fakeATC1.AppendHandlers(
					verifyRegister,
					verifyHeartbeat,
				)
				fakeATC2.AppendHandlers(
					ghttp.CombineHandlers(
						verifyHeartbeat,
						func(w http.ResponseWriter, r *http.Request) { fakeATC2.CloseClientConnections() },
					),
					verifyHeartbeat,
				)
			})

			It("heartbeats faster according to cprInterval", func() {
				Eventually(registrations).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(interval)
				Eventually(heartbeats).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(cprInterval)
				expectedWorker.ActiveContainers = 4
				expectedWorker.ActiveVolumes = 1
				Eventually(heartbeats).Should(Receive(Equal(registration{expectedWorker, 2 * interval})))
			})

			It("goes back to normal after the heartbeat succeeds", func() {
				Eventually(registrations).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(interval)
				Eventually(heartbeats).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(cprInterval)
				Eventually(heartbeats).Should(Receive())

				fakeClock.WaitForWatcherAndIncrement(cprInterval)
				Consistently(heartbeats).ShouldNot(Receive())

				fakeClock.WaitForWatcherAndIncrement(interval - cprInterval)
				expectedWorker.ActiveContainers = 3
				expectedWorker.ActiveVolumes = 0
				Eventually(heartbeats).Should(Receive(Equal(registration{expectedWorker, 2 * interval})))
			})
		})
	})
})
<|endoftext|>


<commit_before>// Copyright 2021 The Prometheus Authors
// 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 chunkenc

import (
	"math/bits"

	"github.com/pkg/errors"
)

// putVarbitInt writes an int64 using varbit encoding with a bit bucketing
// optimized for the dod's observed in histogram buckets, plus a few additional
// buckets for large numbers.
//
// TODO(Dieterbe): We could improve this further: Each branch doesn't need to
// support any values of any of the prior branches. So we can expand the range
// of each branch. Do more with fewer bits. It comes at the price of more
// expensive encoding and decoding (cutting out and later adding back that
// center-piece we skip).
func putVarbitInt(b *bstream, val int64) {
	switch {
	case val == 0: // Precisely 0, needs 1 bit.
		b.writeBit(zero)
	case bitRange(val, 3): // -3 <= val <= 4, needs 5 bits.
		b.writeBits(0b10, 2)
		b.writeBits(uint64(val), 3)
	case bitRange(val, 6): // -31 <= val <= 32, 9 bits.
		b.writeBits(0b110, 3)
		b.writeBits(uint64(val), 6)
	case bitRange(val, 9): // -255 <= val <= 256, 13 bits.
		b.writeBits(0b1110, 4)
		b.writeBits(uint64(val), 9)
	case bitRange(val, 12): // -2047 <= val <= 2048, 17 bits.
		b.writeBits(0b11110, 5)
		b.writeBits(uint64(val), 12)
	case bitRange(val, 18): // -131071 <= val <= 131072, 3 bytes.
		b.writeBits(0b111110, 6)
		b.writeBits(uint64(val), 18)
	case bitRange(val, 25): // -16777215 <= val <= 16777216, 4 bytes.
		b.writeBits(0b1111110, 7)
		b.writeBits(uint64(val), 25)
	case bitRange(val, 56): // -36028797018963967 <= val <= 36028797018963968, 8 bytes.
		b.writeBits(0b11111110, 8)
		b.writeBits(uint64(val), 56)
	default:
		b.writeBits(0b11111111, 8) // Worst case, needs 9 bytes.
		b.writeBits(uint64(val), 64)
	}
}

// readVarbitInt reads an int64 encoced with putVarbitInt.
func readVarbitInt(b *bstreamReader) (int64, error) {
	var d byte
	for i := 0; i < 8; i++ {
		d <<= 1
		bit, err := b.readBitFast()
		if err != nil {
			bit, err = b.readBit()
		}
		if err != nil {
			return 0, err
		}
		if bit == zero {
			break
		}
		d |= 1
	}

	var val int64
	var sz uint8

	switch d {
	case 0b0:
		// val == 0
	case 0b10:
		sz = 3
	case 0b110:
		sz = 6
	case 0b1110:
		sz = 9
	case 0b11110:
		sz = 12
	case 0b111110:
		sz = 18
	case 0b1111110:
		sz = 25
	case 0b11111110:
		sz = 56
	case 0b11111111:
		// Do not use fast because it's very unlikely it will succeed.
		bits, err := b.readBits(64)
		if err != nil {
			return 0, err
		}

		val = int64(bits)
	default:
		return 0, errors.Errorf("invalid bit pattern %b", d)
	}

	if sz != 0 {
		bits, err := b.readBitsFast(sz)
		if err != nil {
			bits, err = b.readBits(sz)
		}
		if err != nil {
			return 0, err
		}
		if bits > (1 << (sz - 1)) {
			// Or something.
			bits = bits - (1 << sz)
		}
		val = int64(bits)
	}

	return val, nil
}

func bitRangeUint(x uint64, nbits int) bool {
	return bits.LeadingZeros64(x) >= 64-nbits
}

// putVarbitUint writes a uint64 using varbit encoding. It uses the same bit
// buckets as putVarbitInt.
func putVarbitUint(b *bstream, val uint64) {
	switch {
	case val == 0: // Precisely 0, needs 1 bit.
		b.writeBit(zero)
	case bitRangeUint(val, 3): // val <= 7, needs 5 bits.
		b.writeBits(0b10, 2)
		b.writeBits(val, 3)
	case bitRangeUint(val, 6): // val <= 63, 9 bits.
		b.writeBits(0b110, 3)
		b.writeBits(val, 6)
	case bitRangeUint(val, 9): // val <= 511, 13 bits.
		b.writeBits(0b1110, 4)
		b.writeBits(val, 9)
	case bitRangeUint(val, 12): // val <= 4095, 17 bits.
		b.writeBits(0b11110, 5)
		b.writeBits(val, 12)
	case bitRangeUint(val, 18): // val <= 262143, 3 bytes.
		b.writeBits(0b111110, 6)
		b.writeBits(val, 18)
	case bitRangeUint(val, 25): // val <= 33554431, 4 bytes.
		b.writeBits(0b1111110, 7)
		b.writeBits(val, 25)
	case bitRangeUint(val, 56): // val <= 72057594037927935, 8 bytes.
		b.writeBits(0b11111110, 8)
		b.writeBits(val, 56)
	default:
		b.writeBits(0b11111111, 8) // Worst case, needs 9 bytes.
		b.writeBits(val, 64)
	}
}

// readVarbitUint reads a uint64 encoced with putVarbitUint.
func readVarbitUint(b *bstreamReader) (uint64, error) {
	var d byte
	for i := 0; i < 8; i++ {
		d <<= 1
		bit, err := b.readBitFast()
		if err != nil {
			bit, err = b.readBit()
		}
		if err != nil {
			return 0, err
		}
		if bit == zero {
			break
		}
		d |= 1
	}

	var (
		bits uint64
		sz   uint8
		err  error
	)

	switch d {
	case 0b0:
		// val == 0
	case 0b10:
		sz = 3
	case 0b110:
		sz = 6
	case 0b1110:
		sz = 9
	case 0b11110:
		sz = 12
	case 0b111110:
		sz = 18
	case 0b1111110:
		sz = 25
	case 0b11111110:
		sz = 56
	case 0b11111111:
		// Do not use fast because it's very unlikely it will succeed.
		bits, err = b.readBits(64)
		if err != nil {
			return 0, err
		}
	default:
		return 0, errors.Errorf("invalid bit pattern %b", d)
	}

	if sz != 0 {
		bits, err = b.readBitsFast(sz)
		if err != nil {
			bits, err = b.readBits(sz)
		}
		if err != nil {
			return 0, err
		}
	}

	return bits, nil
}
<commit_msg>tsdb: Update comment about a possible space optimization (#11303)<commit_after>// Copyright 2021 The Prometheus Authors
// 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 chunkenc

import (
	"math/bits"

	"github.com/pkg/errors"
)

// putVarbitInt writes an int64 using varbit encoding with a bit bucketing
// optimized for the dod's observed in histogram buckets, plus a few additional
// buckets for large numbers.
//
// For optimal space utilization, each branch didn't need to support any values
// of any of the prior branches. So we could expand the range of each branch. Do
// more with fewer bits. It would come at the price of more expensive encoding
// and decoding (cutting out and later adding back that center-piece we
// skip). With the distributions of values we see in practice, we would reduce
// the size by around 1%. A more detailed study would be needed for precise
// values, but it's appears quite certain that we would end up far below 10%,
// which would maybe convince us to invest the increased coding/decoding cost.
func putVarbitInt(b *bstream, val int64) {
	switch {
	case val == 0: // Precisely 0, needs 1 bit.
		b.writeBit(zero)
	case bitRange(val, 3): // -3 <= val <= 4, needs 5 bits.
		b.writeBits(0b10, 2)
		b.writeBits(uint64(val), 3)
	case bitRange(val, 6): // -31 <= val <= 32, 9 bits.
		b.writeBits(0b110, 3)
		b.writeBits(uint64(val), 6)
	case bitRange(val, 9): // -255 <= val <= 256, 13 bits.
		b.writeBits(0b1110, 4)
		b.writeBits(uint64(val), 9)
	case bitRange(val, 12): // -2047 <= val <= 2048, 17 bits.
		b.writeBits(0b11110, 5)
		b.writeBits(uint64(val), 12)
	case bitRange(val, 18): // -131071 <= val <= 131072, 3 bytes.
		b.writeBits(0b111110, 6)
		b.writeBits(uint64(val), 18)
	case bitRange(val, 25): // -16777215 <= val <= 16777216, 4 bytes.
		b.writeBits(0b1111110, 7)
		b.writeBits(uint64(val), 25)
	case bitRange(val, 56): // -36028797018963967 <= val <= 36028797018963968, 8 bytes.
		b.writeBits(0b11111110, 8)
		b.writeBits(uint64(val), 56)
	default:
		b.writeBits(0b11111111, 8) // Worst case, needs 9 bytes.
		b.writeBits(uint64(val), 64)
	}
}

// readVarbitInt reads an int64 encoced with putVarbitInt.
func readVarbitInt(b *bstreamReader) (int64, error) {
	var d byte
	for i := 0; i < 8; i++ {
		d <<= 1
		bit, err := b.readBitFast()
		if err != nil {
			bit, err = b.readBit()
		}
		if err != nil {
			return 0, err
		}
		if bit == zero {
			break
		}
		d |= 1
	}

	var val int64
	var sz uint8

	switch d {
	case 0b0:
		// val == 0
	case 0b10:
		sz = 3
	case 0b110:
		sz = 6
	case 0b1110:
		sz = 9
	case 0b11110:
		sz = 12
	case 0b111110:
		sz = 18
	case 0b1111110:
		sz = 25
	case 0b11111110:
		sz = 56
	case 0b11111111:
		// Do not use fast because it's very unlikely it will succeed.
		bits, err := b.readBits(64)
		if err != nil {
			return 0, err
		}

		val = int64(bits)
	default:
		return 0, errors.Errorf("invalid bit pattern %b", d)
	}

	if sz != 0 {
		bits, err := b.readBitsFast(sz)
		if err != nil {
			bits, err = b.readBits(sz)
		}
		if err != nil {
			return 0, err
		}
		if bits > (1 << (sz - 1)) {
			// Or something.
			bits = bits - (1 << sz)
		}
		val = int64(bits)
	}

	return val, nil
}

func bitRangeUint(x uint64, nbits int) bool {
	return bits.LeadingZeros64(x) >= 64-nbits
}

// putVarbitUint writes a uint64 using varbit encoding. It uses the same bit
// buckets as putVarbitInt.
func putVarbitUint(b *bstream, val uint64) {
	switch {
	case val == 0: // Precisely 0, needs 1 bit.
		b.writeBit(zero)
	case bitRangeUint(val, 3): // val <= 7, needs 5 bits.
		b.writeBits(0b10, 2)
		b.writeBits(val, 3)
	case bitRangeUint(val, 6): // val <= 63, 9 bits.
		b.writeBits(0b110, 3)
		b.writeBits(val, 6)
	case bitRangeUint(val, 9): // val <= 511, 13 bits.
		b.writeBits(0b1110, 4)
		b.writeBits(val, 9)
	case bitRangeUint(val, 12): // val <= 4095, 17 bits.
		b.writeBits(0b11110, 5)
		b.writeBits(val, 12)
	case bitRangeUint(val, 18): // val <= 262143, 3 bytes.
		b.writeBits(0b111110, 6)
		b.writeBits(val, 18)
	case bitRangeUint(val, 25): // val <= 33554431, 4 bytes.
		b.writeBits(0b1111110, 7)
		b.writeBits(val, 25)
	case bitRangeUint(val, 56): // val <= 72057594037927935, 8 bytes.
		b.writeBits(0b11111110, 8)
		b.writeBits(val, 56)
	default:
		b.writeBits(0b11111111, 8) // Worst case, needs 9 bytes.
		b.writeBits(val, 64)
	}
}

// readVarbitUint reads a uint64 encoced with putVarbitUint.
func readVarbitUint(b *bstreamReader) (uint64, error) {
	var d byte
	for i := 0; i < 8; i++ {
		d <<= 1
		bit, err := b.readBitFast()
		if err != nil {
			bit, err = b.readBit()
		}
		if err != nil {
			return 0, err
		}
		if bit == zero {
			break
		}
		d |= 1
	}

	var (
		bits uint64
		sz   uint8
		err  error
	)

	switch d {
	case 0b0:
		// val == 0
	case 0b10:
		sz = 3
	case 0b110:
		sz = 6
	case 0b1110:
		sz = 9
	case 0b11110:
		sz = 12
	case 0b111110:
		sz = 18
	case 0b1111110:
		sz = 25
	case 0b11111110:
		sz = 56
	case 0b11111111:
		// Do not use fast because it's very unlikely it will succeed.
		bits, err = b.readBits(64)
		if err != nil {
			return 0, err
		}
	default:
		return 0, errors.Errorf("invalid bit pattern %b", d)
	}

	if sz != 0 {
		bits, err = b.readBitsFast(sz)
		if err != nil {
			bits, err = b.readBits(sz)
		}
		if err != nil {
			return 0, err
		}
	}

	return bits, nil
}
<|endoftext|>


<commit_before>package ovs

import (
	"fmt"
	"net"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/golang/glog"

	"github.com/openshift/openshift-sdn/plugins/osdn"
	"github.com/openshift/openshift-sdn/plugins/osdn/api"

	kapi "k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/resource"
	kubeletTypes "k8s.io/kubernetes/pkg/kubelet/container"
	knetwork "k8s.io/kubernetes/pkg/kubelet/network"
	utilsets "k8s.io/kubernetes/pkg/util/sets"
)

const (
	SingleTenantPluginName string = "redhat/openshift-ovs-subnet"
	MultiTenantPluginName  string = "redhat/openshift-ovs-multitenant"

	IngressBandwidthAnnotation string = "kubernetes.io/ingress-bandwidth"
	EgressBandwidthAnnotation  string = "kubernetes.io/egress-bandwidth"
	AssignMacVlanAnnotation    string = "pod.network.openshift.io/assign-macvlan"
)

type ovsPlugin struct {
	osdn.OsdnController

	multitenant bool
}

func isPluginMultitenant(pluginName string) bool {
	switch strings.ToLower(pluginName) {
	case MultiTenantPluginName:
		return true
	}
	return false
}

func CreatePlugin(registry *osdn.Registry, pluginName string, hostname string, selfIP string, iptablesSyncPeriod time.Duration) (api.OsdnPlugin, error) {
	plugin := &ovsPlugin{multitenant: isPluginMultitenant(pluginName)}

	err := plugin.BaseInit(registry, plugin, pluginName, hostname, selfIP, iptablesSyncPeriod)
	if err != nil {
		return nil, err
	}

	return plugin, err
}

func (plugin *ovsPlugin) PluginStartMaster(clusterNetwork *net.IPNet, hostSubnetLength uint) error {
	if err := plugin.SubnetStartMaster(clusterNetwork, hostSubnetLength); err != nil {
		return err
	}

	if plugin.multitenant {
		if err := plugin.VnidStartMaster(); err != nil {
			return err
		}
	}

	return nil
}

func (plugin *ovsPlugin) PluginStartNode(mtu uint) error {
	networkChanged, err := plugin.SubnetStartNode(mtu)
	if err != nil {
		return err
	}

	if plugin.multitenant {
		if err := plugin.VnidStartNode(); err != nil {
			return err
		}
	}

	if networkChanged {
		pods, err := plugin.GetLocalPods(kapi.NamespaceAll)
		if err != nil {
			return err
		}
		for _, p := range pods {
			containerID := osdn.GetPodContainerID(&p)
			err = plugin.UpdatePod(p.Namespace, p.Name, kubeletTypes.DockerID(containerID))
			if err != nil {
				glog.Warningf("Could not update pod %q (%s): %s", p.Name, containerID, err)
			}
		}
	}

	return nil
}

//-----------------------------------------------

const (
	setUpCmd    = "setup"
	tearDownCmd = "teardown"
	statusCmd   = "status"
	updateCmd   = "update"
)

func (plugin *ovsPlugin) getExecutable() string {
	return "openshift-sdn-ovs"
}

func (plugin *ovsPlugin) Init(host knetwork.Host) error {
	return nil
}

func (plugin *ovsPlugin) Name() string {
	if plugin.multitenant {
		return MultiTenantPluginName
	} else {
		return SingleTenantPluginName
	}
}

func (plugin *ovsPlugin) Capabilities() utilsets.Int {
	return utilsets.NewInt(knetwork.NET_PLUGIN_CAPABILITY_SHAPING)
}

func (plugin *ovsPlugin) getVNID(namespace string) (string, error) {
	if plugin.multitenant {
		vnid, err := plugin.WaitAndGetVNID(namespace)
		if err != nil {
			return "", err
		}
		return strconv.FormatUint(uint64(vnid), 10), nil
	}

	return "0", nil
}

var minRsrc = resource.MustParse("1k")
var maxRsrc = resource.MustParse("1P")

func parseAndValidateBandwidth(value string) (int64, error) {
	rsrc, err := resource.ParseQuantity(value)
	if err != nil {
		return -1, err
	}

	if rsrc.Value() < minRsrc.Value() {
		return -1, fmt.Errorf("resource value %d is unreasonably small (< %d)", rsrc.Value(), minRsrc.Value())
	}
	if rsrc.Value() > maxRsrc.Value() {
		return -1, fmt.Errorf("resource value %d is unreasonably large (> %d)", rsrc.Value(), maxRsrc.Value())
	}
	return rsrc.Value(), nil
}

func extractBandwidthResources(pod *kapi.Pod) (ingress, egress int64, err error) {
	str, found := pod.Annotations[IngressBandwidthAnnotation]
	if found {
		ingress, err = parseAndValidateBandwidth(str)
		if err != nil {
			return -1, -1, err
		}
	}
	str, found = pod.Annotations[EgressBandwidthAnnotation]
	if found {
		egress, err = parseAndValidateBandwidth(str)
		if err != nil {
			return -1, -1, err
		}
	}
	return ingress, egress, nil
}

func wantsMacvlan(pod *kapi.Pod) (bool, error) {
	val, found := pod.Annotations[AssignMacVlanAnnotation]
	if !found || val != "true" {
		return false, nil
	}
	for _, container := range pod.Spec.Containers {
		if container.SecurityContext.Privileged != nil && *container.SecurityContext.Privileged {
			return true, nil
		}
	}
	return false, fmt.Errorf("Pod has %q annotation but is not privileged", AssignMacVlanAnnotation)
}

func isScriptError(err error) bool {
	_, ok := err.(*exec.ExitError)
	return ok
}

// Get the last command (which is prefixed with "+" because of "set -x") and its output
// (Unless the script ended with "echo ...; exit", in which case we just return the
// echoed text.)
func getScriptError(output []byte) string {
	lines := strings.Split(string(output), "\n")
	last := len(lines)
	for n := last - 1; n >= 0; n-- {
		if strings.HasPrefix(lines[n], "+ exit") {
			last = n
		} else if strings.HasPrefix(lines[n], "+ echo") {
			return strings.Join(lines[n+1:last], "\n")
		} else if strings.HasPrefix(lines[n], "+") {
			return strings.Join(lines[n:], "\n")
		}
	}
	return string(output)
}

func (plugin *ovsPlugin) SetUpPod(namespace string, name string, id kubeletTypes.DockerID) error {
	err := plugin.WaitForPodNetworkReady()
	if err != nil {
		return err
	}

	pod, err := plugin.Registry.GetPod(plugin.HostName, namespace, name)
	if err != nil {
		return err
	}
	if pod == nil {
		return fmt.Errorf("failed to retrieve pod %s/%s", namespace, name)
	}
	ingress, egress, err := extractBandwidthResources(pod)
	if err != nil {
		return fmt.Errorf("failed to parse pod %s/%s ingress/egress quantity: %v", namespace, name, err)
	}
	var ingressStr, egressStr string
	if ingress > 0 {
		ingressStr = fmt.Sprintf("%d", ingress)
	}
	if egress > 0 {
		egressStr = fmt.Sprintf("%d", egress)
	}

	vnidstr, err := plugin.getVNID(namespace)
	if err != nil {
		return err
	}

	macvlan, err := wantsMacvlan(pod)
	if err != nil {
		return err
	}

	out, err := exec.Command(plugin.getExecutable(), setUpCmd, string(id), vnidstr, ingressStr, egressStr, fmt.Sprintf("%t", macvlan)).CombinedOutput()
	glog.V(5).Infof("SetUpPod network plugin output: %s, %v", string(out), err)

	if isScriptError(err) {
		return fmt.Errorf("Error running network setup script: %s", getScriptError(out))
	} else {
		return err
	}
}

func (plugin *ovsPlugin) TearDownPod(namespace string, name string, id kubeletTypes.DockerID) error {
	// The script's teardown functionality doesn't need the VNID
	out, err := exec.Command(plugin.getExecutable(), tearDownCmd, string(id), "-1", "-1", "-1").CombinedOutput()
	glog.V(5).Infof("TearDownPod network plugin output: %s, %v", string(out), err)

	if isScriptError(err) {
		return fmt.Errorf("Error running network teardown script: %s", getScriptError(out))
	} else {
		return err
	}
}

func (plugin *ovsPlugin) Status(namespace string, name string, id kubeletTypes.DockerID) (*knetwork.PodNetworkStatus, error) {
	return nil, nil
}

func (plugin *ovsPlugin) UpdatePod(namespace string, name string, id kubeletTypes.DockerID) error {
	vnidstr, err := plugin.getVNID(namespace)
	if err != nil {
		return err
	}

	out, err := exec.Command(plugin.getExecutable(), updateCmd, string(id), vnidstr).CombinedOutput()
	glog.V(5).Infof("UpdatePod network plugin output: %s, %v", string(out), err)

	if isScriptError(err) {
		return fmt.Errorf("Error running network update script: %s", getScriptError(out))
	} else {
		return err
	}
}

func (plugin *ovsPlugin) Event(name string, details map[string]interface{}) {
}
<commit_msg>Added IsOpenShiftNetworkPlugin and IsOpenShiftMultitenantNetworkPlugin methods<commit_after>package ovs

import (
	"fmt"
	"net"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/golang/glog"

	"github.com/openshift/openshift-sdn/plugins/osdn"
	"github.com/openshift/openshift-sdn/plugins/osdn/api"

	kapi "k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/resource"
	kubeletTypes "k8s.io/kubernetes/pkg/kubelet/container"
	knetwork "k8s.io/kubernetes/pkg/kubelet/network"
	utilsets "k8s.io/kubernetes/pkg/util/sets"
)

const (
	SingleTenantPluginName string = "redhat/openshift-ovs-subnet"
	MultiTenantPluginName  string = "redhat/openshift-ovs-multitenant"

	IngressBandwidthAnnotation string = "kubernetes.io/ingress-bandwidth"
	EgressBandwidthAnnotation  string = "kubernetes.io/egress-bandwidth"
	AssignMacVlanAnnotation    string = "pod.network.openshift.io/assign-macvlan"
)

type ovsPlugin struct {
	osdn.OsdnController

	multitenant bool
}

func IsOpenShiftNetworkPlugin(pluginName string) bool {
	switch strings.ToLower(pluginName) {
	case SingleTenantPluginName, MultiTenantPluginName:
		return true
	}
	return false
}

func IsOpenShiftMultitenantNetworkPlugin(pluginName string) bool {
	if strings.ToLower(pluginName) == MultiTenantPluginName {
		return true
	}
	return false
}

func CreatePlugin(registry *osdn.Registry, pluginName string, hostname string, selfIP string, iptablesSyncPeriod time.Duration) (api.OsdnPlugin, error) {
	plugin := &ovsPlugin{multitenant: IsOpenShiftMultitenantNetworkPlugin(pluginName)}

	err := plugin.BaseInit(registry, plugin, pluginName, hostname, selfIP, iptablesSyncPeriod)
	if err != nil {
		return nil, err
	}

	return plugin, err
}

func (plugin *ovsPlugin) PluginStartMaster(clusterNetwork *net.IPNet, hostSubnetLength uint) error {
	if err := plugin.SubnetStartMaster(clusterNetwork, hostSubnetLength); err != nil {
		return err
	}

	if plugin.multitenant {
		if err := plugin.VnidStartMaster(); err != nil {
			return err
		}
	}

	return nil
}

func (plugin *ovsPlugin) PluginStartNode(mtu uint) error {
	networkChanged, err := plugin.SubnetStartNode(mtu)
	if err != nil {
		return err
	}

	if plugin.multitenant {
		if err := plugin.VnidStartNode(); err != nil {
			return err
		}
	}

	if networkChanged {
		pods, err := plugin.GetLocalPods(kapi.NamespaceAll)
		if err != nil {
			return err
		}
		for _, p := range pods {
			containerID := osdn.GetPodContainerID(&p)
			err = plugin.UpdatePod(p.Namespace, p.Name, kubeletTypes.DockerID(containerID))
			if err != nil {
				glog.Warningf("Could not update pod %q (%s): %s", p.Name, containerID, err)
			}
		}
	}

	return nil
}

//-----------------------------------------------

const (
	setUpCmd    = "setup"
	tearDownCmd = "teardown"
	statusCmd   = "status"
	updateCmd   = "update"
)

func (plugin *ovsPlugin) getExecutable() string {
	return "openshift-sdn-ovs"
}

func (plugin *ovsPlugin) Init(host knetwork.Host) error {
	return nil
}

func (plugin *ovsPlugin) Name() string {
	if plugin.multitenant {
		return MultiTenantPluginName
	} else {
		return SingleTenantPluginName
	}
}

func (plugin *ovsPlugin) Capabilities() utilsets.Int {
	return utilsets.NewInt(knetwork.NET_PLUGIN_CAPABILITY_SHAPING)
}

func (plugin *ovsPlugin) getVNID(namespace string) (string, error) {
	if plugin.multitenant {
		vnid, err := plugin.WaitAndGetVNID(namespace)
		if err != nil {
			return "", err
		}
		return strconv.FormatUint(uint64(vnid), 10), nil
	}

	return "0", nil
}

var minRsrc = resource.MustParse("1k")
var maxRsrc = resource.MustParse("1P")

func parseAndValidateBandwidth(value string) (int64, error) {
	rsrc, err := resource.ParseQuantity(value)
	if err != nil {
		return -1, err
	}

	if rsrc.Value() < minRsrc.Value() {
		return -1, fmt.Errorf("resource value %d is unreasonably small (< %d)", rsrc.Value(), minRsrc.Value())
	}
	if rsrc.Value() > maxRsrc.Value() {
		return -1, fmt.Errorf("resource value %d is unreasonably large (> %d)", rsrc.Value(), maxRsrc.Value())
	}
	return rsrc.Value(), nil
}

func extractBandwidthResources(pod *kapi.Pod) (ingress, egress int64, err error) {
	str, found := pod.Annotations[IngressBandwidthAnnotation]
	if found {
		ingress, err = parseAndValidateBandwidth(str)
		if err != nil {
			return -1, -1, err
		}
	}
	str, found = pod.Annotations[EgressBandwidthAnnotation]
	if found {
		egress, err = parseAndValidateBandwidth(str)
		if err != nil {
			return -1, -1, err
		}
	}
	return ingress, egress, nil
}

func wantsMacvlan(pod *kapi.Pod) (bool, error) {
	val, found := pod.Annotations[AssignMacVlanAnnotation]
	if !found || val != "true" {
		return false, nil
	}
	for _, container := range pod.Spec.Containers {
		if container.SecurityContext.Privileged != nil && *container.SecurityContext.Privileged {
			return true, nil
		}
	}
	return false, fmt.Errorf("Pod has %q annotation but is not privileged", AssignMacVlanAnnotation)
}

func isScriptError(err error) bool {
	_, ok := err.(*exec.ExitError)
	return ok
}

// Get the last command (which is prefixed with "+" because of "set -x") and its output
// (Unless the script ended with "echo ...; exit", in which case we just return the
// echoed text.)
func getScriptError(output []byte) string {
	lines := strings.Split(string(output), "\n")
	last := len(lines)
	for n := last - 1; n >= 0; n-- {
		if strings.HasPrefix(lines[n], "+ exit") {
			last = n
		} else if strings.HasPrefix(lines[n], "+ echo") {
			return strings.Join(lines[n+1:last], "\n")
		} else if strings.HasPrefix(lines[n], "+") {
			return strings.Join(lines[n:], "\n")
		}
	}
	return string(output)
}

func (plugin *ovsPlugin) SetUpPod(namespace string, name string, id kubeletTypes.DockerID) error {
	err := plugin.WaitForPodNetworkReady()
	if err != nil {
		return err
	}

	pod, err := plugin.Registry.GetPod(plugin.HostName, namespace, name)
	if err != nil {
		return err
	}
	if pod == nil {
		return fmt.Errorf("failed to retrieve pod %s/%s", namespace, name)
	}
	ingress, egress, err := extractBandwidthResources(pod)
	if err != nil {
		return fmt.Errorf("failed to parse pod %s/%s ingress/egress quantity: %v", namespace, name, err)
	}
	var ingressStr, egressStr string
	if ingress > 0 {
		ingressStr = fmt.Sprintf("%d", ingress)
	}
	if egress > 0 {
		egressStr = fmt.Sprintf("%d", egress)
	}

	vnidstr, err := plugin.getVNID(namespace)
	if err != nil {
		return err
	}

	macvlan, err := wantsMacvlan(pod)
	if err != nil {
		return err
	}

	out, err := exec.Command(plugin.getExecutable(), setUpCmd, string(id), vnidstr, ingressStr, egressStr, fmt.Sprintf("%t", macvlan)).CombinedOutput()
	glog.V(5).Infof("SetUpPod network plugin output: %s, %v", string(out), err)

	if isScriptError(err) {
		return fmt.Errorf("Error running network setup script: %s", getScriptError(out))
	} else {
		return err
	}
}

func (plugin *ovsPlugin) TearDownPod(namespace string, name string, id kubeletTypes.DockerID) error {
	// The script's teardown functionality doesn't need the VNID
	out, err := exec.Command(plugin.getExecutable(), tearDownCmd, string(id), "-1", "-1", "-1").CombinedOutput()
	glog.V(5).Infof("TearDownPod network plugin output: %s, %v", string(out), err)

	if isScriptError(err) {
		return fmt.Errorf("Error running network teardown script: %s", getScriptError(out))
	} else {
		return err
	}
}

func (plugin *ovsPlugin) Status(namespace string, name string, id kubeletTypes.DockerID) (*knetwork.PodNetworkStatus, error) {
	return nil, nil
}

func (plugin *ovsPlugin) UpdatePod(namespace string, name string, id kubeletTypes.DockerID) error {
	vnidstr, err := plugin.getVNID(namespace)
	if err != nil {
		return err
	}

	out, err := exec.Command(plugin.getExecutable(), updateCmd, string(id), vnidstr).CombinedOutput()
	glog.V(5).Infof("UpdatePod network plugin output: %s, %v", string(out), err)

	if isScriptError(err) {
		return fmt.Errorf("Error running network update script: %s", getScriptError(out))
	} else {
		return err
	}
}

func (plugin *ovsPlugin) Event(name string, details map[string]interface{}) {
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"github.com/gruntwork-io/terragrunt/options"
	"github.com/gruntwork-io/terragrunt/util"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"
	"testing"
)

func TestAlreadyHaveLatestCodeLocalFilePath(t *testing.T) {
	t.Parallel()

	canonicalUrl := fmt.Sprintf("file://%s", absPath(t, "../test/fixture-download-source/hello-world"))
	downloadDir := "does-not-exist"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, false)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirDoesNotExist(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com"
	downloadDir := "does-not-exist"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, false)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirExistsNoVersionNoVersionFile(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com"
	downloadDir := "../test/fixture-download-source/download-dir-empty"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, false)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirExistsNoVersionWithVersionFile(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com"
	downloadDir := "../test/fixture-download-source/download-dir-version-file-no-query"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, true)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirExistsWithVersionNoVersionFile(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com?ref=v0.0.1"
	downloadDir := "../test/fixture-download-source/download-dir-empty"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, false)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirExistsWithVersionAndVersionFile(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com?ref=v0.0.1"
	downloadDir := "../test/fixture-download-source/download-dir-version-file"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, true)
}

func TestDownloadTerraformSourceIfNecessaryLocalDirToEmptyDir(t *testing.T) {
	t.Parallel()

	canonicalUrl := fmt.Sprintf("file://%s", absPath(t, "../test/fixture-download-source/hello-world"))
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World")
}

func TestDownloadTerraformSourceIfNecessaryLocalDirToAlreadyDownloadedDir(t *testing.T) {
	t.Parallel()

	canonicalUrl := fmt.Sprintf("file://%s", absPath(t, "../test/fixture-download-source/hello-world"))
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-2", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlToEmptyDir(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlToAlreadyDownloadedDir(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-2", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World 2")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlToAlreadyDownloadedDirDifferentVersion(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world?ref=v0.9.7"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-2", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlToAlreadyDownloadedDirSameVersion(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world?ref=v0.9.7"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-version-remote", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World version remote")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlOverrideSource(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world?ref=v0.9.7"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-version-remote", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, true, "# Hello, World")
}

func testDownloadTerraformSourceIfNecessary(t *testing.T, canonicalUrl string, downloadDir string, sourceUpdate bool, expectedFileContents string) {
	terraformSource := &TerraformSource{
		CanonicalSourceURL: parseUrl(t, canonicalUrl),
		DownloadDir:        downloadDir,
		WorkingDir:         downloadDir,
		VersionFile:        util.JoinPath(downloadDir, "version-file.txt"),
	}

	terragruntOptions := options.NewTerragruntOptionsForTest("./should-not-be-used")
	terragruntOptions.SourceUpdate = sourceUpdate

	err := downloadTerraformSourceIfNecessary(terraformSource, terragruntOptions)
	assert.Nil(t, err, "For terraform source %v: %v", terraformSource, err)

	expectedFilePath := util.JoinPath(downloadDir, "main.tf")
	if assert.True(t, util.FileExists(expectedFilePath), "For terraform source %v", terraformSource) {
		actualFileContents := readFile(t, expectedFilePath)
		assert.Equal(t, expectedFileContents, actualFileContents, "For terraform source %v", terraformSource)
	}
}

func testAlreadyHaveLatestCode(t *testing.T, canonicalUrl string, downloadDir string, expected bool) {
	terraformSource := &TerraformSource{
		CanonicalSourceURL: parseUrl(t, canonicalUrl),
		DownloadDir:        downloadDir,
		WorkingDir:         downloadDir,
		VersionFile:        util.JoinPath(downloadDir, "version-file.txt"),
	}

	actual, err := alreadyHaveLatestCode(terraformSource)
	assert.Nil(t, err, "Unexpected error for terraform source %v: %v", terraformSource, err)
	assert.Equal(t, expected, actual, "For terraform source %v", terraformSource)
}

func tmpDir(t *testing.T) string {
	dir, err := ioutil.TempDir("", "download-source-test")
	if err != nil {
		t.Fatal(err)
	}
	return dir
}

func absPath(t *testing.T, path string) string {
	abs, err := filepath.Abs(path)
	if err != nil {
		t.Fatal(err)
	}
	return abs
}

func parseUrl(t *testing.T, rawUrl string) *url.URL {
	parsed, err := url.Parse(rawUrl)
	if err != nil {
		t.Fatal(err)
	}
	return parsed
}

func readFile(t *testing.T, path string) string {
	contents, err := util.ReadFileAsString(path)
	if err != nil {
		t.Fatal(err)
	}
	return contents
}

func copyFolder(t *testing.T, src string, dest string) {
	err := util.CopyFolderContents(src, dest)
	if err != nil {
		t.Fatal(err)
	}
}
<commit_msg>Fix compile error in download_source_test<commit_after>package cli

import (
	"fmt"
	"github.com/gruntwork-io/terragrunt/options"
	"github.com/gruntwork-io/terragrunt/util"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"
	"testing"
	"github.com/gruntwork-io/terragrunt/config"
)

func TestAlreadyHaveLatestCodeLocalFilePath(t *testing.T) {
	t.Parallel()

	canonicalUrl := fmt.Sprintf("file://%s", absPath(t, "../test/fixture-download-source/hello-world"))
	downloadDir := "does-not-exist"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, false)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirDoesNotExist(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com"
	downloadDir := "does-not-exist"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, false)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirExistsNoVersionNoVersionFile(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com"
	downloadDir := "../test/fixture-download-source/download-dir-empty"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, false)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirExistsNoVersionWithVersionFile(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com"
	downloadDir := "../test/fixture-download-source/download-dir-version-file-no-query"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, true)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirExistsWithVersionNoVersionFile(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com?ref=v0.0.1"
	downloadDir := "../test/fixture-download-source/download-dir-empty"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, false)
}

func TestAlreadyHaveLatestCodeRemoteFilePathDownloadDirExistsWithVersionAndVersionFile(t *testing.T) {
	t.Parallel()

	canonicalUrl := "http://www.some-url.com?ref=v0.0.1"
	downloadDir := "../test/fixture-download-source/download-dir-version-file"

	testAlreadyHaveLatestCode(t, canonicalUrl, downloadDir, true)
}

func TestDownloadTerraformSourceIfNecessaryLocalDirToEmptyDir(t *testing.T) {
	t.Parallel()

	canonicalUrl := fmt.Sprintf("file://%s", absPath(t, "../test/fixture-download-source/hello-world"))
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World")
}

func TestDownloadTerraformSourceIfNecessaryLocalDirToAlreadyDownloadedDir(t *testing.T) {
	t.Parallel()

	canonicalUrl := fmt.Sprintf("file://%s", absPath(t, "../test/fixture-download-source/hello-world"))
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-2", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlToEmptyDir(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlToAlreadyDownloadedDir(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-2", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World 2")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlToAlreadyDownloadedDirDifferentVersion(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world?ref=v0.9.7"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-2", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlToAlreadyDownloadedDirSameVersion(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world?ref=v0.9.7"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-version-remote", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, false, "# Hello, World version remote")
}

func TestDownloadTerraformSourceIfNecessaryRemoteUrlOverrideSource(t *testing.T) {
	t.Parallel()

	canonicalUrl := "github.com/gruntwork-io/terragrunt//test/fixture-download-source/hello-world?ref=v0.9.7"
	downloadDir := tmpDir(t)
	defer os.Remove(downloadDir)

	copyFolder(t, "../test/fixture-download-source/hello-world-version-remote", downloadDir)

	testDownloadTerraformSourceIfNecessary(t, canonicalUrl, downloadDir, true, "# Hello, World")
}

func testDownloadTerraformSourceIfNecessary(t *testing.T, canonicalUrl string, downloadDir string, sourceUpdate bool, expectedFileContents string) {
	terraformSource := &TerraformSource{
		CanonicalSourceURL: parseUrl(t, canonicalUrl),
		DownloadDir:        downloadDir,
		WorkingDir:         downloadDir,
		VersionFile:        util.JoinPath(downloadDir, "version-file.txt"),
	}

	terragruntOptions := options.NewTerragruntOptionsForTest("./should-not-be-used")
	terragruntOptions.SourceUpdate = sourceUpdate
	terragruntConfig := config.TerragruntConfig{}

	err := downloadTerraformSourceIfNecessary(terraformSource, &terragruntConfig, terragruntOptions)
	assert.Nil(t, err, "For terraform source %v: %v", terraformSource, err)

	expectedFilePath := util.JoinPath(downloadDir, "main.tf")
	if assert.True(t, util.FileExists(expectedFilePath), "For terraform source %v", terraformSource) {
		actualFileContents := readFile(t, expectedFilePath)
		assert.Equal(t, expectedFileContents, actualFileContents, "For terraform source %v", terraformSource)
	}
}

func testAlreadyHaveLatestCode(t *testing.T, canonicalUrl string, downloadDir string, expected bool) {
	terraformSource := &TerraformSource{
		CanonicalSourceURL: parseUrl(t, canonicalUrl),
		DownloadDir:        downloadDir,
		WorkingDir:         downloadDir,
		VersionFile:        util.JoinPath(downloadDir, "version-file.txt"),
	}

	actual, err := alreadyHaveLatestCode(terraformSource)
	assert.Nil(t, err, "Unexpected error for terraform source %v: %v", terraformSource, err)
	assert.Equal(t, expected, actual, "For terraform source %v", terraformSource)
}

func tmpDir(t *testing.T) string {
	dir, err := ioutil.TempDir("", "download-source-test")
	if err != nil {
		t.Fatal(err)
	}
	return dir
}

func absPath(t *testing.T, path string) string {
	abs, err := filepath.Abs(path)
	if err != nil {
		t.Fatal(err)
	}
	return abs
}

func parseUrl(t *testing.T, rawUrl string) *url.URL {
	parsed, err := url.Parse(rawUrl)
	if err != nil {
		t.Fatal(err)
	}
	return parsed
}

func readFile(t *testing.T, path string) string {
	contents, err := util.ReadFileAsString(path)
	if err != nil {
		t.Fatal(err)
	}
	return contents
}

func copyFolder(t *testing.T, src string, dest string) {
	err := util.CopyFolderContents(src, dest)
	if err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import "flag"

var RustReserved = []string{
	/*"break",        "default",      "func",         "interface",    "select",
	"case",         "defer",        "go",           "map",          "struct",
	"chan",         "else",         "goto",         "package",      "switch",
	"const",        "fallthrough",  "if",           "range",        "type",
	"continue",     "for",          "import",       "return",       "var",
	"bool",			"byte", 		"len", 			"open", 		"file",
	"close", 		"load", 		"copy",			"new",*/
}

//This is the Java compiler for uct.
var Rust bool

func init() {
	flag.BoolVar(&Rust, "rs", false, "Target Rust")

	RegisterAssembler(RustAssembly, &Rust, "rs", "//")

	for _, word := range RustReserved {
		RustAssembly[word] = Reserved()
	}
}

var RustAssembly = Assemblable{
	//Special commands.
	"HEADER": Instruction{
		Data: `
#![allow(warnings)]
extern crate num;
extern crate rand;
mod stack;

use num::bigint::BigInt;
use num::bigint::{ToBigInt, Sign};
use num::ToPrimitive;

use stack::Stack;
use stack::Div;
use stack::Mod;
use stack::Pipe;
use stack::Pointer;
`,
		Args: 1,
	},

	"FOOTER": Instruction{},

	"FILE": Instruction{
		Path: "stack.rs",
	},

	"RUST": Instruction{All:true},

	"NUMBER": is("%s.to_bigint().unwrap()", 1),
	"BIG": 	is("BigInt(`%s`)", 1),
	"SIZE":   is("stack.Borrow[%s.address].len().to_bigint().unwrap()", 1),
	"STRING": is("stack.NewStringArray(%s)", 1),
	"ERRORS":  is("stack.ERROR", 1),

	//Experimental.
	//"EVAL": is(`stack.Eval[stack.Grab().String()](stack)`),
	//"EVALUATION": is(`stack.Eval["%s"] = %s`, 2),

	"SOFTWARE": Instruction{
		Data:   "fn main() { let mut stack = Stack::new();",
		Indent: 1,
	},
	"EXIT": Instruction{
		Indented:    1,
		Data:        "}\n",
		Indent:      -1,
		Indentation: -1,
		Else: &Instruction{
			Data: "exit",
		},
	},

	"FUNCTION": is("fn %s(mut stack: &mut Stack) {", 1, 1),
	"RETURN": Instruction{
		Indented:    1,
		Data:        "}\n",
		Indent:      -1,
		Indentation: -1,
		Else: &Instruction{
			Data: "return",
		},
	},

	"SCOPE": is(`stack.Relay(Pipe{Function:%s});`, 1),

	"EXE": is("%s.Exe(&mut stack);", 1),

	"PUSH": is("stack.Push(%s.clone());", 1),
	"PULL": is("let mut %s = stack.Pull();", 1),

	"PUT":   is("stack.Put(%s);", 1),
	"POP":   is("let mut %s = stack.Pop();", 1),
	"PLACE": is("stack.ActiveArray.address = %s.address;", 1),
	"ARRAY":  is("let mut %s = stack.Array();", 1),
	"MAKE":  is("stack.Make();"),
	"RENAME": is("%s = stack.ActiveArray.clone(); stack.Library[stack.ActiveArray.address].set(stack.Library[stack.ActiveArray.address].get()+1);", 1),
	"RELOAD": is("%s = stack.Take();", 1),

	"SHARE": is("stack.Share(%s.clone());", 1),
	"GRAB":  is("let mut %s = stack.Grab();", 1),

	"PIPE": is("let mut %s = stack.Pipe();", 1),
	"RELAY": is("stack.Relay(%s);", 1),
	"TAKE":  is("let mut %s = stack.Take();", 1),

	"GET": is("let mut %s = stack.Get();", 1),
	"SET": is("stack.Set(%s);", 1),

	"VAR": is("let mut %s = 0.to_bigint().unwrap();", 1),

	//"INBOX":   is("stack.Share(<- stack.Inbox)"),
	//"OUTBOX":   is("stack.Outbox <- stack.Grab()"),

	"OPEN":   is("stack.Open();"),
	"DELETE": is("stack.Delete();"),
	"MOVE":   is("stack.Move();"),
	"EXECUTE":is("stack.Execute();"),
	"LINK":   is("stack.Link();"),
	"CONNECT":is("stack.Connect();"),
	"SLICE":  is("stack.Slice();"),
	"LOAD":   is("stack.Load();"),
	"OUT":    is("stack.Out();"),
	"STAT":   is("stack.Info();"),
	"IN":     is("stack.In();"),
	"STDOUT": is("stack.Stdout();"),
	"STDIN":  is("stack.Stdin();"),
	"HEAP":  is("stack.Heap();"),
	"HEAPIT":  is("stack.HeapIt();"),

	"CLOSE": is("%s.Close();", 1),

	"LOOP":   is("loop {", 0, 1),
	"BREAK":  is("break;"),
	"REPEAT": is("}", 0, -1, -1),

	"IF":   is("if %s != 0.to_bigint().unwrap() {", 1, 1),
	"ELSE": is("} else {", 0, 0, -1),
	"END":  is("}", 0, -1, -1),

	"RUN":  is("%s(&mut stack);", 1),

	"DATA": is(" ", 2),

	"FORK": is("go %s(stack.Copy())\n", 1),

	"ADD": is("%s = &%s + &%s;", 3),
	"SUB": is("%s = &%s - &%s;", 3),
	"MUL": is("%s = &%s * &%s;", 3),
	"DIV": is("%s = Div(&%s, &%s);", 3),
	"MOD": is("%s = Mod(&%s, &%s);", 3),
	"POW": is("%s = num::pow(%s, %s.to_usize().unwrap());", 3),

	"SLT": is("%s = if %s < %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SEQ": is("%s = if %s == %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SGE": is("%s = if %s >= %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SGT": is("%s = if %s > %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SNE": is("%s = if %s != %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SLE": is("%s = if %s <= %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),

	"JOIN": is("%s = stack.Join(&%s, &%s);", 3),
	"ERROR": is("stack.ERROR = %s;", 1),
}
<commit_msg>Add rust reserved words.<commit_after>package main

import "flag"

var RustReserved = []string{
	"abstract",	"alignof",	"as",		"become",	"box",
	"break",	"const",	"continue",	"crate",	"do",
	"else",		"enum",		"extern",	"false",	"final",
	"fn",		"for",		"if",		"impl",		"in",
	"let",		"loop",		"macro",	"match",	"mod",
	"move",		"mut",		"offsetof",	"override",	"priv",
	"proc",		"pub",		"pure",		"ref",		"return",
	"Self",		"self",		"sizeof",	"static",	"struct",
	"super",	"trait",	"true",		"type",		"typeof",
	"unsafe",	"unsized",	"use",		"virtual",	"where",
	"while",	"yield",
}

//This is the Java compiler for uct.
var Rust bool

func init() {
	flag.BoolVar(&Rust, "rs", false, "Target Rust")

	RegisterAssembler(RustAssembly, &Rust, "rs", "//")

	for _, word := range RustReserved {
		RustAssembly[word] = Reserved()
	}
}

var RustAssembly = Assemblable{
	//Special commands.
	"HEADER": Instruction{
		Data: `
#![allow(warnings)]
extern crate num;
extern crate rand;
mod stack;

use num::bigint::BigInt;
use num::bigint::{ToBigInt, Sign};
use num::ToPrimitive;

use stack::Stack;
use stack::Div;
use stack::Mod;
use stack::Pipe;
use stack::Pointer;
`,
		Args: 1,
	},

	"FOOTER": Instruction{},

	"FILE": Instruction{
		Path: "stack.rs",
	},

	"RUST": Instruction{All:true},

	"NUMBER": is("%s.to_bigint().unwrap()", 1),
	"BIG": 	is("BigInt(`%s`)", 1),
	"SIZE":   is("stack.Borrow[%s.address].len().to_bigint().unwrap()", 1),
	"STRING": is("stack.NewStringArray(%s)", 1),
	"ERRORS":  is("stack.ERROR", 1),

	//Experimental.
	//"EVAL": is(`stack.Eval[stack.Grab().String()](stack)`),
	//"EVALUATION": is(`stack.Eval["%s"] = %s`, 2),

	"SOFTWARE": Instruction{
		Data:   "fn main() { let mut stack = Stack::new();",
		Indent: 1,
	},
	"EXIT": Instruction{
		Indented:    1,
		Data:        "}\n",
		Indent:      -1,
		Indentation: -1,
		Else: &Instruction{
			Data: "exit",
		},
	},

	"FUNCTION": is("fn %s(mut stack: &mut Stack) {", 1, 1),
	"RETURN": Instruction{
		Indented:    1,
		Data:        "}\n",
		Indent:      -1,
		Indentation: -1,
		Else: &Instruction{
			Data: "return",
		},
	},

	"SCOPE": is(`stack.Relay(Pipe{Function:%s});`, 1),

	"EXE": is("%s.Exe(&mut stack);", 1),

	"PUSH": is("stack.Push(%s.clone());", 1),
	"PULL": is("let mut %s = stack.Pull();", 1),

	"PUT":   is("stack.Put(%s);", 1),
	"POP":   is("let mut %s = stack.Pop();", 1),
	"PLACE": is("stack.ActiveArray.address = %s.address;", 1),
	"ARRAY":  is("let mut %s = stack.Array();", 1),
	"MAKE":  is("stack.Make();"),
	"RENAME": is("%s = stack.ActiveArray.clone(); stack.Library[stack.ActiveArray.address].set(stack.Library[stack.ActiveArray.address].get()+1);", 1),
	"RELOAD": is("%s = stack.Take();", 1),

	"SHARE": is("stack.Share(%s.clone());", 1),
	"GRAB":  is("let mut %s = stack.Grab();", 1),

	"PIPE": is("let mut %s = stack.Pipe();", 1),
	"RELAY": is("stack.Relay(%s);", 1),
	"TAKE":  is("let mut %s = stack.Take();", 1),

	"GET": is("let mut %s = stack.Get();", 1),
	"SET": is("stack.Set(%s);", 1),

	"VAR": is("let mut %s = 0.to_bigint().unwrap();", 1),

	//"INBOX":   is("stack.Share(<- stack.Inbox)"),
	//"OUTBOX":   is("stack.Outbox <- stack.Grab()"),

	"OPEN":   is("stack.Open();"),
	"DELETE": is("stack.Delete();"),
	"MOVE":   is("stack.Move();"),
	"EXECUTE":is("stack.Execute();"),
	"LINK":   is("stack.Link();"),
	"CONNECT":is("stack.Connect();"),
	"SLICE":  is("stack.Slice();"),
	"LOAD":   is("stack.Load();"),
	"OUT":    is("stack.Out();"),
	"STAT":   is("stack.Info();"),
	"IN":     is("stack.In();"),
	"STDOUT": is("stack.Stdout();"),
	"STDIN":  is("stack.Stdin();"),
	"HEAP":  is("stack.Heap();"),
	"HEAPIT":  is("stack.HeapIt();"),

	"CLOSE": is("%s.Close();", 1),

	"LOOP":   is("loop {", 0, 1),
	"BREAK":  is("break;"),
	"REPEAT": is("}", 0, -1, -1),

	"IF":   is("if %s != 0.to_bigint().unwrap() {", 1, 1),
	"ELSE": is("} else {", 0, 0, -1),
	"END":  is("}", 0, -1, -1),

	"RUN":  is("%s(&mut stack);", 1),

	"DATA": is(" ", 2),

	"FORK": is("go %s(stack.Copy())\n", 1),

	"ADD": is("%s = &%s + &%s;", 3),
	"SUB": is("%s = &%s - &%s;", 3),
	"MUL": is("%s = &%s * &%s;", 3),
	"DIV": is("%s = Div(&%s, &%s);", 3),
	"MOD": is("%s = Mod(&%s, &%s);", 3),
	"POW": is("%s = num::pow(%s, %s.to_usize().unwrap());", 3),

	"SLT": is("%s = if %s < %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SEQ": is("%s = if %s == %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SGE": is("%s = if %s >= %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SGT": is("%s = if %s > %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SNE": is("%s = if %s != %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),
	"SLE": is("%s = if %s <= %s { 1.to_bigint().unwrap() } else { 0.to_bigint().unwrap() };", 3),

	"JOIN": is("%s = stack.Join(&%s, &%s);", 3),
	"ERROR": is("stack.ERROR = %s;", 1),
}
<|endoftext|>


<commit_before>package main

import (
	"dns"
	"log"
)

func serve(w dns.ResponseWriter, req *dns.Msg, z *dns.Zone) {
	log.Printf("incoming %s %s\n", req.Question[0].Name, dns.Rr_str[req.Question[0].Qtype])
	// Check for referral
	// if we find something with NonAuth = true, it means
	// we need to return referral
	nss := z.Predecessor(req.Question[0].Name)
	if nss != nil && nss.NonAuth {
		log.Printf("Referral")
		m := new(dns.Msg)
		m.SetReply(req)
		m.Ns = nss.RR[dns.TypeNS]
		// lookup the a records for additional, only when
		// in baliwick
		w.Write(m)
		return

	}

	// For now:
	// look up name -> yes, continue, no -> nxdomain
	node := z.Find(req.Question[0].Name)
	if node == nil {
		m := new(dns.Msg)
		m.SetRcode(req, dns.RcodeNameError)
		w.Write(m)
		return
	}
	apex := z.Find(z.Origin)

	// Name found, look for type, yes, answer, no
	if rrs, ok := node.RR[req.Question[0].Qtype]; ok {
		// Need to look at class to but... no
		m := new(dns.Msg)
		m.SetReply(req)
		m.MsgHdr.Authoritative = true
		m.Answer = rrs
		// auth section
		m.Ns = apex.RR[dns.TypeNS]
		w.Write(m)
		return
	} else {
		// nodate reply
		// soa in auth section
		m := new(dns.Msg)
		m.SetReply(req)
		m.Ns = apex.RR[dns.TypeSOA]
		w.Write(m)
		return
	}
	// See RFC 1035...
	m := new(dns.Msg)
	m.SetRcode(req, dns.RcodeNameError)

	w.Write(m)
}
<commit_msg>some fixes<commit_after>package main

import (
	"dns"
	"log"
)

func serve(w dns.ResponseWriter, req *dns.Msg, z *dns.Zone) {
	log.Printf("incoming %s %s\n", req.Question[0].Name, dns.Rr_str[req.Question[0].Qtype])
	// Check for referral
	// if we find something with NonAuth = true, it means
	// we need to return referral
	nss := z.Predecessor(req.Question[0].Name)
	if nss != nil && nss.NonAuth {
		log.Printf("Referral")
		m := new(dns.Msg)
		m.SetReply(req)
		m.Ns = nss.RR[dns.TypeNS]
		// lookup the a records for additional, only when
		// in baliwick
		w.Write(m)
		return

	}

	// If we don't have the name return NXDOMAIN
	node := z.Find(req.Question[0].Name)
	if node == nil {
		m := new(dns.Msg)
		m.SetRcode(req, dns.RcodeNameError)
		w.Write(m)
		return
	}
	apex := z.Find(z.Origin)

	// Check if we have the type
	if rrs, ok := node.RR[req.Question[0].Qtype]; ok {
		m := new(dns.Msg)
		m.SetReply(req)
		m.MsgHdr.Authoritative = true
		m.Answer = rrs
		m.Ns = apex.RR[dns.TypeNS] // Add auth section
		w.Write(m)
		return
	} else {
		// We don't have the type, check if we have non-auth ns servers, if so
		// we need to return a referral
		if nss, ok := node.RR[dns.TypeNS]; ok && node.NonAuth {
			log.Printf("Also Referral!!")
			m := new(dns.Msg)
			m.SetReply(req)
			m.Ns = nss
			// lookup the a records for additional, only when
			// in baliwick
			w.Write(m)

		}
		m := new(dns.Msg) // Nodata reply
		m.SetReply(req)
		m.Ns = apex.RR[dns.TypeSOA]
		w.Write(m)
		return
	}
	// See RFC 1035...
	m := new(dns.Msg)
	m.SetRcode(req, dns.RcodeNameError)

	w.Write(m)
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"os"

	"github.com/minio/cli"
	"github.com/tinylib/msgp/msgp"
)

func main() {
	app := cli.NewApp()
	app.Copyright = "MinIO, Inc."
	app.Usage = "xl.meta to JSON"
	app.Version = "0.0.1"
	app.HideHelpCommand = true

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Usage: "path to xl.meta file",
			Name:  "f, file",
		},
	}

	app.Action = func(c *cli.Context) error {
		r, err := os.Open(c.String("file"))
		if err != nil {
			return err
		}
		r.Seek(8, io.SeekStart)
		defer r.Close()
		_, err = msgp.CopyToJSON(os.Stdout, r)
		return err
	}
	app.Run(os.Args)
}
<commit_msg>Add neater metadata printer (#9898)<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"

	"github.com/minio/cli"
	"github.com/tinylib/msgp/msgp"
)

var xlHeader = [4]byte{'X', 'L', '2', ' '}

func main() {
	app := cli.NewApp()
	app.Copyright = "MinIO, Inc."
	app.Usage = "xl.meta to JSON"
	app.Version = "0.0.1"
	app.HideHelpCommand = true

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Usage: "Print each file as a separate line without formatting",
			Name:  "ndjson",
		},
	}

	app.Action = func(c *cli.Context) error {
		for _, file := range c.Args() {
			var r io.Reader
			switch file {
			case "-":
				r = os.Stdin
			default:
				f, err := os.Open(file)
				if err != nil {
					return err
				}
				defer f.Close()
				r = f
			}

			// Read header
			var tmp [4]byte
			_, err := io.ReadFull(r, tmp[:])
			if err != nil {
				return err
			}
			if !bytes.Equal(tmp[:], xlHeader[:]) {
				return fmt.Errorf("xlMeta: unknown XLv2 header, expected %v, got %v", xlHeader[:4], tmp[:4])
			}
			// Skip version check for now
			_, err = io.ReadFull(r, tmp[:])
			if err != nil {
				return err
			}

			var buf bytes.Buffer
			_, err = msgp.CopyToJSON(&buf, r)
			if err != nil {
				return err
			}
			if c.Bool("ndjson") {
				fmt.Println(buf.String())
				continue
			}
			var msi map[string]interface{}
			dec := json.NewDecoder(&buf)
			// Use number to preserve integers.
			dec.UseNumber()
			err = dec.Decode(&msi)
			if err != nil {
				return err
			}
			b, err := json.MarshalIndent(msi, "", "  ")
			if err != nil {
				return err
			}
			fmt.Println(string(b))
		}
		return nil
	}
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package tcproxy

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/pkg/errors"
	schematypes "github.com/taskcluster/go-schematypes"
	"github.com/taskcluster/taskcluster-worker/engines"
	"github.com/taskcluster/taskcluster-worker/plugins"
	"github.com/taskcluster/taskcluster-worker/runtime"
)

type provider struct {
	plugins.PluginProviderBase
}

type plugin struct {
	plugins.PluginBase
}

type taskPlugin struct {
	plugins.TaskPluginBase
	monitor runtime.Monitor
	context *runtime.TaskContext
}

func init() {
	plugins.Register("tcproxy", provider{})
}

func (provider) NewPlugin(options plugins.PluginOptions) (plugins.Plugin, error) {
	return &plugin{}, nil
}

func (p *plugin) PayloadSchema() schematypes.Object {
	return payloadSchema
}

func (p *plugin) NewTaskPlugin(options plugins.TaskPluginOptions) (plugins.TaskPlugin, error) {
	var P payload
	schematypes.MustValidateAndMap(payloadSchema, options.Payload, &P)

	// If disabled we return nothing
	if P.DisableTaskclusterProxy {
		return plugins.TaskPluginBase{}, nil
	}

	return &taskPlugin{
		monitor: options.Monitor,
		context: options.TaskContext,
	}, nil
}

func (p *taskPlugin) BuildSandbox(sandboxBuilder engines.SandboxBuilder) error {
	err := sandboxBuilder.AttachProxy("tcproxy", p)
	if err == engines.ErrFeatureNotSupported {
		// Fire off a warning, and then do nothing...
		p.monitor.ReportWarning(err, "plugin 'tcproxy' is enabled, but the engine doesn't support proxy attachments")
		return nil
	}
	if err == engines.ErrNamingConflict {
		return runtime.NewMalformedPayloadError("the proxy name 'tcproxy' is already in use")
	}
	if _, ok := runtime.IsMalformedPayloadError(err); ok {
		// the name "tcproxy" is not allowed by the engine, we assume it to be safe,
		// so if it's not we'll panic
		panic(errors.Wrap(err, "proxy name 'tcproxy' is not permitted by the engine"))
	}
	return nil
}

func (p *taskPlugin) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Parse request URL
	raw := r.URL.Path
	if strings.HasPrefix(raw, "/") {
		raw = raw[1:]
	}
	u, err := url.Parse("https://" + raw)
	if err != nil {
		debug("bad URL: '%s'", r.URL.Path)
		p.context.LogError(fmt.Sprintf("tcproxy received path: '%s' which it failed to parse as a URL", raw))
		w.WriteHeader(http.StatusBadRequest)
		data, _ := json.MarshalIndent(struct {
			Code    string `json:"code"`
			Message string `json:"message"`
		}{
			Code:    "InvalidRequestUrl",
			Message: "tcproxy assumes <path> is on the form <hostname>/<resource>[?<query>], instead received: '" + raw + "'",
		}, "", "  ")
		w.Write(data)
		return
	}

	debug("proxing: '%s'", raw)

	// Create request
	req, _ := http.NewRequest(r.Method, u.String(), r.Body)

	// Add headers except 'Host'
	r.Header.Del("Host")
	for k, v := range r.Header {
		req.Header[k] = v
	}

	// Add signature
	signature, err := p.context.Authorizer().SignHeader(r.Method, u, nil)
	if err != nil {
		incidentID := p.monitor.ReportError(
			errors.Wrap(err, "SignHeader failed"),
			"SignHeader failed for URL: ", u.String(),
		)
		p.context.LogError(fmt.Sprintf("tcproxy expirenced an internal error, incidentID: %s", incidentID))
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(`{
      "code": "InternalServerError",
      "message": "internal error in taskcluster-worker proxy, incidentId: ` + incidentID + `"
    }`))
		return
	}
	req.Header.Set("Authorization", signature)

	// Set TaskContext, so request can be aborted
	req.WithContext(p.context)

	// Send request
	res, err := http.DefaultClient.Do(req)

	// Handle task aborting
	if p.context.Err() != nil {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(`{
      "code": "TaskResolved",
      "message": "task was resolved, aborting proxied requests"
    }`))
		// Quick clean-up body, if for some reason cancel happened after response
		if err == nil && res.Body != nil {
			res.Body.Close()
		}
		return
	}

	// Handle connection errors
	if err != nil {
		// This could be task that canceled the request, or task connection broke
		// and r.Body returned an error or something like that... Or it could be
		// the request URL was bad, or service was down...
		p.context.Log(fmt.Sprintf(
			"tcproxy was unable to forward request to %s, error: %s",
			u.String(), err,
		))
		w.WriteHeader(http.StatusInternalServerError)
		data, _ := json.MarshalIndent(struct {
			Code    string `json:"code"`
			Message string `json:"message"`
		}{
			Code:    "ProxyRequestFailed",
			Message: "tcproxy failed forwarding request to '" + u.String() + "'",
		}, "", "  ")
		w.Write(data)
		return
	}

	// Set headers from res
	for k, v := range res.Header {
		w.Header()[k] = v
	}

	// Write header
	w.WriteHeader(res.StatusCode)

	// Copy body
	_, err = io.Copy(w, res.Body)
	if err != nil {
		debug("failed for proxy response, error: %s", err)
		p.context.Log(fmt.Sprintf("tcproxy failed to proxy the entire response from: %s", u.String()))
	}
}
<commit_msg>Fixes caught by linter<commit_after>package tcproxy

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/pkg/errors"
	schematypes "github.com/taskcluster/go-schematypes"
	"github.com/taskcluster/taskcluster-worker/engines"
	"github.com/taskcluster/taskcluster-worker/plugins"
	"github.com/taskcluster/taskcluster-worker/runtime"
)

type provider struct {
	plugins.PluginProviderBase
}

type plugin struct {
	plugins.PluginBase
}

type taskPlugin struct {
	plugins.TaskPluginBase
	monitor runtime.Monitor
	context *runtime.TaskContext
}

func init() {
	plugins.Register("tcproxy", provider{})
}

func (provider) NewPlugin(options plugins.PluginOptions) (plugins.Plugin, error) {
	return &plugin{}, nil
}

func (p *plugin) PayloadSchema() schematypes.Object {
	return payloadSchema
}

func (p *plugin) NewTaskPlugin(options plugins.TaskPluginOptions) (plugins.TaskPlugin, error) {
	var P payload
	schematypes.MustValidateAndMap(payloadSchema, options.Payload, &P)

	// If disabled we return nothing
	if P.DisableTaskclusterProxy {
		return plugins.TaskPluginBase{}, nil
	}

	return &taskPlugin{
		monitor: options.Monitor,
		context: options.TaskContext,
	}, nil
}

func (p *taskPlugin) BuildSandbox(sandboxBuilder engines.SandboxBuilder) error {
	err := sandboxBuilder.AttachProxy("tcproxy", p)
	if err == engines.ErrFeatureNotSupported {
		// Fire off a warning, and then do nothing...
		p.monitor.ReportWarning(err, "plugin 'tcproxy' is enabled, but the engine doesn't support proxy attachments")
		return nil
	}
	if err == engines.ErrNamingConflict {
		return runtime.NewMalformedPayloadError("the proxy name 'tcproxy' is already in use")
	}
	if _, ok := runtime.IsMalformedPayloadError(err); ok {
		// the name "tcproxy" is not allowed by the engine, we assume it to be safe,
		// so if it's not we'll panic
		panic(errors.Wrap(err, "proxy name 'tcproxy' is not permitted by the engine"))
	}
	return nil
}

func (p *taskPlugin) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Parse request URL
	raw := strings.TrimPrefix(r.URL.Path, "/")
	u, err := url.Parse("https://" + raw)
	if err != nil {
		debug("bad URL: '%s'", r.URL.Path)
		p.context.LogError(fmt.Sprintf("tcproxy received path: '%s' which it failed to parse as a URL", raw))
		w.WriteHeader(http.StatusBadRequest)
		data, _ := json.MarshalIndent(struct {
			Code    string `json:"code"`
			Message string `json:"message"`
		}{
			Code:    "InvalidRequestUrl",
			Message: "tcproxy assumes <path> is on the form <hostname>/<resource>[?<query>], instead received: '" + raw + "'",
		}, "", "  ")
		w.Write(data)
		return
	}

	debug("proxing: '%s'", raw)

	// Create request
	req, _ := http.NewRequest(r.Method, u.String(), r.Body)

	// Add headers except 'Host'
	r.Header.Del("Host")
	for k, v := range r.Header {
		req.Header[k] = v
	}

	// Add signature
	signature, err := p.context.Authorizer().SignHeader(r.Method, u, nil)
	if err != nil {
		incidentID := p.monitor.ReportError(
			errors.Wrap(err, "SignHeader failed"),
			"SignHeader failed for URL: ", u.String(),
		)
		p.context.LogError(fmt.Sprintf("tcproxy expirenced an internal error, incidentID: %s", incidentID))
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(`{
      "code": "InternalServerError",
      "message": "internal error in taskcluster-worker proxy, incidentId: ` + incidentID + `"
    }`))
		return
	}
	req.Header.Set("Authorization", signature)

	// Set TaskContext, so request can be aborted
	req = req.WithContext(p.context)

	// Send request
	res, err := http.DefaultClient.Do(req)

	// Handle task aborting
	if p.context.Err() != nil {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(`{
      "code": "TaskResolved",
      "message": "task was resolved, aborting proxied requests"
    }`))
		// Quick clean-up body, if for some reason cancel happened after response
		if err == nil && res.Body != nil {
			res.Body.Close()
		}
		return
	}

	// Handle connection errors
	if err != nil {
		// This could be task that canceled the request, or task connection broke
		// and r.Body returned an error or something like that... Or it could be
		// the request URL was bad, or service was down...
		p.context.Log(fmt.Sprintf(
			"tcproxy was unable to forward request to %s, error: %s",
			u.String(), err,
		))
		w.WriteHeader(http.StatusInternalServerError)
		data, _ := json.MarshalIndent(struct {
			Code    string `json:"code"`
			Message string `json:"message"`
		}{
			Code:    "ProxyRequestFailed",
			Message: "tcproxy failed forwarding request to '" + u.String() + "'",
		}, "", "  ")
		w.Write(data)
		return
	}

	// Set headers from res
	for k, v := range res.Header {
		w.Header()[k] = v
	}

	// Write header
	w.WriteHeader(res.StatusCode)

	// Copy body
	_, err = io.Copy(w, res.Body)
	if err != nil {
		debug("failed for proxy response, error: %s", err)
		p.context.Log(fmt.Sprintf("tcproxy failed to proxy the entire response from: %s", u.String()))
	}
}
<|endoftext|>


<commit_before>package twitter

import (
	"fmt"

	"go.uber.org/zap"

	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
	"github.com/jirwin/quadlek/quadlek"
)

func load(consumerKey, consumerSecret, accessToken, accessSecret string, filter map[string]string) func(bot *quadlek.Bot, store *quadlek.Store) error {

	return func(bot *quadlek.Bot, store *quadlek.Store) error {
		go func() {
			config := oauth1.NewConfig(consumerKey, consumerSecret)
			token := oauth1.NewToken(accessToken, accessSecret)
			httpClient := config.Client(oauth1.NoContext, token)
			client := twitter.NewClient(httpClient)

			followFilters := []string{}
			for follow := range filter {
				followFilters = append(followFilters, follow)
			}

			filterParams := &twitter.StreamFilterParams{
				Follow:        followFilters,
				StallWarnings: twitter.Bool(true),
			}

			stream, err := client.Streams.Filter(filterParams)
			if err != nil {
				zap.L().Error("Error streaming tweets.", zap.Error(err))
				return
			}

			for msg := range stream.Messages {
				switch m := msg.(type) {
				case *twitter.Tweet:
					if channel, ok := filter[m.User.IDStr]; ok {
						if m.InReplyToStatusIDStr != "" {
							zap.L().Info("skipping reply tweet", zap.Any("tweet", m))
							continue
						}
						if m.RetweetedStatus != nil {
							zap.L().Info("Got a tweet containing a retweet", zap.Any("tweet", m))

							if replyChannel, ok := filter[m.RetweetedStatus.User.IDStr]; ok && channel == replyChannel {
								zap.L().Info("Tweet contains retweet from already monitored account, cancelling message", zap.Any("tweet", m))
								continue
							}
						}
						twitterUrl := fmt.Sprintf("https://nitter.quadlek.dev/%s/status/%s", m.User.ScreenName, m.IDStr)
						chanId, err := bot.GetChannelId(channel)
						if err != nil {
							zap.L().Error("unable to find channel.", zap.Error(err))
							continue
						}
						bot.Say(chanId, twitterUrl)
					}
				}
			}
		}()

		return nil
	}
}

func Register(consumerKey, consumerSecret, accessToken, accessSecret string, filter map[string]string) quadlek.Plugin {
	return quadlek.MakePlugin(
		"quadlek-twitter",
		nil,
		nil,
		nil,
		nil,
		load(consumerKey, consumerSecret, accessToken, accessSecret, filter),
	)
}
<commit_msg>Use twitter instead of nitter (#28)<commit_after>package twitter

import (
	"fmt"

	"go.uber.org/zap"

	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
	"github.com/jirwin/quadlek/quadlek"
)

func load(consumerKey, consumerSecret, accessToken, accessSecret string, filter map[string]string) func(bot *quadlek.Bot, store *quadlek.Store) error {

	return func(bot *quadlek.Bot, store *quadlek.Store) error {
		go func() {
			config := oauth1.NewConfig(consumerKey, consumerSecret)
			token := oauth1.NewToken(accessToken, accessSecret)
			httpClient := config.Client(oauth1.NoContext, token)
			client := twitter.NewClient(httpClient)

			followFilters := []string{}
			for follow := range filter {
				followFilters = append(followFilters, follow)
			}

			filterParams := &twitter.StreamFilterParams{
				Follow:        followFilters,
				StallWarnings: twitter.Bool(true),
			}

			stream, err := client.Streams.Filter(filterParams)
			if err != nil {
				zap.L().Error("Error streaming tweets.", zap.Error(err))
				return
			}

			for msg := range stream.Messages {
				switch m := msg.(type) {
				case *twitter.Tweet:
					if channel, ok := filter[m.User.IDStr]; ok {
						if m.InReplyToStatusIDStr != "" {
							zap.L().Info("skipping reply tweet", zap.Any("tweet", m))
							continue
						}
						if m.RetweetedStatus != nil {
							zap.L().Info("Got a tweet containing a retweet", zap.Any("tweet", m))

							if replyChannel, ok := filter[m.RetweetedStatus.User.IDStr]; ok && channel == replyChannel {
								zap.L().Info("Tweet contains retweet from already monitored account, cancelling message", zap.Any("tweet", m))
								continue
							}
						}
						twitterUrl := fmt.Sprintf("https://twitter.com/%s/status/%s", m.User.ScreenName, m.IDStr)
						chanId, err := bot.GetChannelId(channel)
						if err != nil {
							zap.L().Error("unable to find channel.", zap.Error(err))
							continue
						}
						bot.Say(chanId, twitterUrl)
					}
				}
			}
		}()

		return nil
	}
}

func Register(consumerKey, consumerSecret, accessToken, accessSecret string, filter map[string]string) quadlek.Plugin {
	return quadlek.MakePlugin(
		"quadlek-twitter",
		nil,
		nil,
		nil,
		nil,
		load(consumerKey, consumerSecret, accessToken, accessSecret, filter),
	)
}
<|endoftext|>