github.com/gramework/gramework@v1.8.1-0.20231027140105-82555c9057f5/mw/xhostname/xhostname.go (about)

     1  // Package xhostname middleware provides `X-Hostname` header in each request
     2  // and useful when using scalable container platform to see
     3  // which host sent you current response.
     4  //
     5  // To use this package, just call `xhostname.Setup`:
     6  //
     7  //		app := gramework.New()
     8  //		xhostname.Setup(app)
     9  package xhostname
    10  
    11  import (
    12  	"os"
    13  
    14  	"github.com/gramework/gramework"
    15  )
    16  
    17  const (
    18  	// HeaderKey is the header name which we use
    19  	HeaderKey = "X-Hostname"
    20  )
    21  
    22  var hostname string
    23  
    24  func init() {
    25  	// kubernetes provides a HOSTNAME in pods
    26  	if h := os.Getenv("HOSTNAME"); len(h) > 0 {
    27  		hostname = h
    28  		gramework.Logger.WithField("hostname", h).Info("using HOSTNAME env as a hostname")
    29  		return
    30  	}
    31  	var err error
    32  	hostname, err = os.Hostname()
    33  	if err != nil {
    34  		gramework.Logger.WithError(err).Error("could not get hostname")
    35  		return
    36  	}
    37  	gramework.Logger.WithField("hostname", hostname).Info("using os.Hostname env as a hostname")
    38  }
    39  
    40  // Setup registers middleware in the provided app
    41  func Setup(app *gramework.App) {
    42  	err := app.UseAfterRequest(serveXHost)
    43  	if err != nil {
    44  		app.Logger.WithError(err).WithField("package", "mw/xhostname").Error("could not register middleware")
    45  	}
    46  }
    47  
    48  func serveXHost(ctx *gramework.Context) {
    49  	ctx.Response.Header.Add(HeaderKey, hostname)
    50  }