Golang http.MaxBytesReader函数代码示例

// Read the request body into the provided variable, it will read the body as json
// if the specified "v" isn't ([]byte, io.Writer) .
func (self Context) GetBody(v interface{}, max int64) (err error) {
	switch v.(type) {
	case []byte:
		_, err = http.MaxBytesReader(self.Res, self.Req.Body, max).Read(v.([]byte))
	case io.Writer:
		_, err = io.Copy(v.(io.Writer), http.MaxBytesReader(self.Res, self.Req.Body, max))
	default:
		err = json.NewDecoder(http.MaxBytesReader(self.Res, self.Req.Body, max)).Decode(v)
	}
	return err
}

func main() {

	var handlerInstance *HandlerObject = NewHandlerObject(123)

	var httpMux *http.ServeMux = http.NewServeMux()
	httpServer := &http.Server{
		Addr:           ":8080",
		Handler:        httpMux,
		ReadTimeout:    2 * time.Second,
		WriteTimeout:   2 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	var httpsMux *http.ServeMux = http.NewServeMux()
	httpsServer := &http.Server{
		Addr:           ":8443",
		Handler:        httpsMux,
		ReadTimeout:    10 * time.Second,
		WriteTimeout:   10 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	httpMux.Handle("/foo", handlerInstance)
	httpMux.HandleFunc("/bar", func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%+v", r)
		r.Body = http.MaxBytesReader(w, nopCloser{r.Body}, 1024)
		fmt.Fprintf(w, "HandlerFunc, %q", html.EscapeString(r.URL.Path))
	})

	httpsMux.Handle("/foo", handlerInstance)
	httpsMux.HandleFunc("/bar", func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%+v", r)
		r.Body = http.MaxBytesReader(w, nopCloser{r.Body}, 65536)
		fmt.Fprintf(w, "HandlerFunc, %q", html.EscapeString(r.URL.Path))
	})

	go func() {
		log.Println("Before starting HTTPS listener...")
		err := httpsServer.ListenAndServeTLS("server.crt", "server.key.insecure")
		if err != nil {
			log.Fatal("HTTPS listener couldn't start")
		}
	}()

	log.Println("Before starting HTTP listener...")
	err := httpServer.ListenAndServe()
	if err != nil {
		log.Fatal("HTTP listener couldn't start")
	}
}

func dataHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Header().Set("Allow", "POST")
		w.WriteHeader(http.StatusMethodNotAllowed)
		w.Write([]byte("You must send a POST request to get data."))
		return
	}

	r.Body = http.MaxBytesReader(w, r.Body, MAX_REQSIZE)
	payload, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.Write([]byte(fmt.Sprintf("Could not read received POST payload: %v", err)))
		return
	}

	var wrapper RespWrapper = RespWrapper{w}

	uuidBytes, startTime, endTime, pw, token, _, success := parseDataRequest(string(payload), wrapper)

	if success {
		var loginsession *LoginSession
		if token != "" {
			loginsession = validateToken(token)
			if loginsession == nil {
				w.Write([]byte(ERROR_INVALID_TOKEN))
				return
			}
		}
		if hasPermission(loginsession, uuidBytes) {
			dr.MakeDataRequest(uuidBytes, startTime, endTime, uint8(pw), wrapper)
		} else {
			wrapper.GetWriter().Write([]byte("[]"))
		}
	}
}

func newContext(w http.ResponseWriter, r *http.Request) (*Context, error) {
	var err error

	c := new(Context)

	c.Request = r
	c.RequestId = atomic.AddInt64(&globalReqId, 1)

	if w != nil {
		c.ResponseHeader = w.Header()
	}
	c.ResponseWriter = w

	c.Values = make(map[string]interface{})

	if r.Body != nil {
		mr := http.MaxBytesReader(w, r.Body, MaxBodyLength)

		c.RawPostData, err = ioutil.ReadAll(mr)
		r.Body.Close()

		if err != nil {
			return c, NewError(err.Error(), StatusBadRequest)
		}
	}

	return c, nil
}

func (h *httpHandler) handlePost(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
	content, err := getContentFromForm(r)
	size := int64(len(content))
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if err := h.stats.MakeSpaceFor(size); err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
	}
	id, err := h.store.Put(content)
	if err != nil {
		log.Printf("Unknown error on POST: %v", err)
		h.stats.FreeSpace(size)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	storage.SetupPasteDeletion(h.store, h.stats, id, size, *lifeTime)
	url := fmt.Sprintf("%s/%s", *siteURL, id)
	switch r.URL.Path {
	case "/redirect":
		http.Redirect(w, r, url, 302)
	default:
		fmt.Fprintln(w, url)
	}
}

func TestRequestBodyLimit(t *testing.T) {
	const limit = 1 << 20
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.Body = http.MaxBytesReader(w, r.Body, limit)
		n, err := io.Copy(ioutil.Discard, r.Body)
		if err == nil {
			t.Errorf("expected error from io.Copy")
		}
		if n != limit {
			t.Errorf("io.Copy = %d, want %d", n, limit)
		}
	}))
	defer ts.Close()

	nWritten := int64(0)
	req, _ := http.NewRequest("POST", ts.URL, io.LimitReader(countReader{neverEnding('a'), &nWritten}, limit*200))

	// Send the POST, but don't care it succeeds or not.  The
	// remote side is going to reply and then close the TCP
	// connection, and HTTP doesn't really define if that's
	// allowed or not.  Some HTTP clients will get the response
	// and some (like ours, currently) will complain that the
	// request write failed, without reading the response.
	//
	// But that's okay, since what we're really testing is that
	// the remote side hung up on us before we wrote too much.
	_, _ = http.DefaultClient.Do(req)

	if nWritten > limit*100 {
		t.Errorf("handler restricted the request body to %d bytes, but client managed to write %d",
			limit, nWritten)
	}
}

func Urlresolve(conn *irc.Connection, resp string, message string) {
	a := strings.Split(message, " ")
	for _, b := range a {
		if strings.HasPrefix(b, "http://") || strings.HasPrefix(b, "https://") {
			response, err := http.Get(b)
			if err != nil {
				fmt.Println(err)
				return
			} else {
				defer response.Body.Close()
				c := http.MaxBytesReader(nil, response.Body, 10000)
				d := new(bytes.Buffer)
				d.ReadFrom(c)
				e := d.String()
				f := strings.Split(string(e), "<")
				for _, g := range f {
					if strings.Contains(g, "title>") ||
						strings.Contains(g, "Title>") ||
						strings.Contains(g, "TITLE>") {
						h := strings.TrimSpace(g)
						i := strings.Split(h, ">")
						j := i[1]
						k := strings.Replace(j, "\n", "", -1)
						l := strings.TrimSpace(k)
						conn.Privmsgf(resp, l)
						break
					}
				}
			}
		}
	}
}

func handleBuild(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	var insecure bool
	rawInsecure := r.FormValue("insecure")
	if rawInsecure != "" && rawInsecure == "true" {
		insecure = true
	}
	j := &job{
		insecure: insecure,
		pkg:      r.URL.Path,
		tar:      http.MaxBytesReader(w, r.Body, MaxTarSize),
		done:     make(chan struct{}),
	}
	Q <- j
	<-j.done
	const httpTooLarge = "http: request body too large"
	if j.err != nil && j.err.Error() == httpTooLarge {
		http.Error(w, httpTooLarge, http.StatusRequestEntityTooLarge)
		return
	}
	if j.err != nil {
		log.Println(j.err)
		http.Error(w, "unprocessable entity", 422)
		w.Write(j.out)
		return
	}
	defer j.bin.Close()
	http.ServeContent(w, r, "", time.Time{}, j.bin)
}

func UploadFile(path string, maxMemory int64, w http.ResponseWriter, r *http.Request) (string, error) {
	r.Body = http.MaxBytesReader(w, r.Body, maxMemory)

	file, _, err := r.FormFile("file")
	if err != nil {
		return "", err
	}
	defer file.Close()

	randName, err := str.Rand(32)
	if err != nil {
		return "", err
	}

	filePath := filepath.Join(path, randName)

	f, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE, 0666)
	if err != nil {
		return "", err
	}
	defer f.Close()

	_, err = io.Copy(f, file)

	return randName, err
}

func readReqJSON(w http.ResponseWriter, r *http.Request, n int64, v interface{}) bool {
	err := json.NewDecoder(http.MaxBytesReader(w, r.Body, n)).Decode(v)
	if err != nil {
		http.Error(w, "unprocessable entity", 422)
	}
	return err == nil
}

// CreateC handles the multipart form upload and creates an encrypted file
func CreateC(c *gin.Context) {
	var err error
	var duration time.Duration
	var once bool

	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, conf.C.SizeLimit*utils.MegaByte)

	once = c.PostForm("once") != ""
	d := c.DefaultPostForm("duration", "1d")

	if val, ok := models.DurationMap[d]; ok {
		duration = val
	} else {
		logger.ErrC(c, "server", "Invalid duration", d)
		c.String(http.StatusBadRequest, "Invalid duration\n")
		c.AbortWithStatus(http.StatusBadRequest)
		return
	}

	fd, h, err := c.Request.FormFile("file")
	if err != nil {
		logger.ErrC(c, "server", "Couldn't read file", err)
		c.String(http.StatusRequestEntityTooLarge, "Entity is too large (Max : %v MB)\n", conf.C.SizeLimit)
		c.AbortWithStatus(http.StatusRequestEntityTooLarge)
		return
	}
	defer fd.Close()

	res := models.NewResourceFromForm(h, once, duration)
	k, err := res.WriteEncrypted(fd)
	if err != nil {
		logger.ErrC(c, "server", "Couldn't write file", err)
		c.String(http.StatusInternalServerError, "Something went wrong on the server. Try again later.")
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}

	if conf.C.DiskQuota > 0 {
		if models.S.CurrentSize+uint64(res.Size) > uint64(conf.C.DiskQuota*utils.GigaByte) {
			logger.ErrC(c, "server", "Quota exceeded")
			c.String(http.StatusBadRequest, "Insufficient disk space. Try again later.")
			c.AbortWithStatus(http.StatusBadRequest)
			os.Remove(path.Join(conf.C.UploadDir, res.Key))
			return
		}
	}

	if err = res.Save(); err != nil {
		logger.ErrC(c, "server", "Couldn't save in the database", err)
		c.String(http.StatusInternalServerError, "Something went wrong on the server. Try again later.")
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	res.LogCreated(c)
	ns := conf.C.NameServer
	if conf.C.AppendPort {
		ns = fmt.Sprintf("%s:%d", conf.C.NameServer, conf.C.Port)
	}
	c.String(http.StatusCreated, "%v://%s/v/%s/%s\n", utils.DetectScheme(c), ns, res.Key, k)
}

func runHandler(resp http.ResponseWriter, req *http.Request,
	fn func(resp http.ResponseWriter, req *http.Request) error, errfn httputil.Error) {
	defer func() {
		if rv := recover(); rv != nil {
			err := errors.New("handler panic")
			logError(req, err, rv)
			errfn(resp, req, http.StatusInternalServerError, err)
		}
	}()

	if s := req.Header.Get("X-Real-Ip"); s != "" && httputil.StripPort(req.RemoteAddr) == "127.0.0.1" {
		req.RemoteAddr = s
	}

	req.Body = http.MaxBytesReader(resp, req.Body, 2048)
	req.ParseForm()
	var rb httputil.ResponseBuffer
	err := fn(&rb, req)
	if err == nil {
		rb.WriteTo(resp)
	} else if e, ok := err.(*httpError); ok {
		if e.status >= 500 {
			logError(req, err, nil)
		}
		errfn(resp, req, e.status, e.err)
	} else {
		logError(req, err, nil)
		errfn(resp, req, http.StatusInternalServerError, err)
	}
}

// Feed the body of req into the OR port, and write any data read from the OR
// port back to w.
func transact(session *Session, w http.ResponseWriter, req *http.Request) error {
	body := http.MaxBytesReader(w, req.Body, maxPayloadLength+1)
	_, err := io.Copy(session.Or, body)
	if err != nil {
		return fmt.Errorf("error copying body to ORPort: %s", scrubError(err))
	}

	buf := make([]byte, maxPayloadLength)
	session.Or.SetReadDeadline(time.Now().Add(turnaroundTimeout))
	n, err := session.Or.Read(buf)
	if err != nil {
		if e, ok := err.(net.Error); !ok || !e.Timeout() {
			httpInternalServerError(w)
			// Don't scrub err here because it always refers to localhost.
			return fmt.Errorf("reading from ORPort: %s", err)
		}
	}
	// log.Printf("read %d bytes from ORPort: %q", n, buf[:n])
	// Set a Content-Type to prevent Go and the CDN from trying to guess.
	w.Header().Set("Content-Type", "application/octet-stream")
	n, err = w.Write(buf[:n])
	if err != nil {
		return fmt.Errorf("error writing to response: %s", scrubError(err))
	}
	// log.Printf("wrote %d bytes to response", n)
	return nil
}

// This method handles all requests.  It dispatches to handleInternal after
// handling / adapting websocket connections.
func handle(w http.ResponseWriter, r *http.Request) {
	if maxRequestSize := int64(Config.IntDefault("http.maxrequestsize", 0)); maxRequestSize > 0 {
		r.Body = http.MaxBytesReader(w, r.Body, maxRequestSize)
	}

	upgrade := r.Header.Get("Upgrade")
	if upgrade == "websocket" || upgrade == "Websocket" {
		websocket.Handler(func(ws *websocket.Conn) {
			//Override default Read/Write timeout with sane value for a web socket request
			ws.SetDeadline(time.Now().Add(time.Hour * 24))
			r.Method = "WS"
			handleInternal(w, r, ws)
		}).ServeHTTP(w, r)
	} else {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
		w.Header().Set("Access-Control-Allow-Headers",
			"Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, X-session, X-adm-token")

		// Stop here for a Preflighted OPTIONS request.
		if r.Method == "OPTIONS" {
			return
		}

		handleInternal(w, r, nil)
	}
}

// convertHTTPRequestToAPIRequest converts the HTTP request query
// parameters and request body to the JSON object import format
// expected by the API request handlers.
func convertHTTPRequestToAPIRequest(
	w http.ResponseWriter,
	r *http.Request,
	requestBodyName string) (requestJSONObject, error) {

	params := make(requestJSONObject)

	for name, values := range r.URL.Query() {
		for _, value := range values {
			params[name] = value
			// Note: multiple values per name are ignored
			break
		}
	}

	if requestBodyName != "" {
		r.Body = http.MaxBytesReader(w, r.Body, MAX_API_PARAMS_SIZE)
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			return nil, psiphon.ContextError(err)
		}
		var bodyParams requestJSONObject
		err = json.Unmarshal(body, &bodyParams)
		if err != nil {
			return nil, psiphon.ContextError(err)
		}
		params[requestBodyName] = bodyParams
	}

	return params, nil
}