/*
 * Copyright (c) 2000, 2013, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Oracle designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Oracle in the LICENSE file that accompanied this code.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

/*
 *
 *  (C) Copyright IBM Corp. 1999 All Rights Reserved.
 *  Copyright 1997 The Open Group Research Institute.  All rights reserved.
 */

package sun.security.krb5.internal;

import java.io.*;
import java.net.*;
import sun.security.util.IOUtils;

public abstract class NetClient implements AutoCloseable {
    public static NetClient getInstance(String protocol, String hostname, int port,
            int timeout) throws IOException {
        if (protocol.equals("TCP")) {
            return new TCPClient(hostname, port, timeout);
        } else {
            return new UDPClient(hostname, port, timeout);
        }
    }

    abstract public void send(byte[] data) throws IOException;
    abstract public byte[] receive() throws IOException;
    abstract public void close() throws IOException;
}

class TCPClient extends NetClient {

    private Socket tcpSocket;
    private BufferedOutputStream out;
    private BufferedInputStream in;

    TCPClient(String hostname, int port, int timeout)
            throws IOException {
        tcpSocket = new Socket();
        tcpSocket.connect(new InetSocketAddress(hostname, port), timeout);
        out = new BufferedOutputStream(tcpSocket.getOutputStream());
        in = new BufferedInputStream(tcpSocket.getInputStream());
        tcpSocket.setSoTimeout(timeout);
    }

    @Override
    public void send(byte[] data) throws IOException {
        byte[] lenField = new byte[4];
        intToNetworkByteOrder(data.length, lenField, 0, 4);
        out.write(lenField);

        out.write(data);
        out.flush();
    }

    @Override
    public byte[] receive() throws IOException {
        byte[] lenField = new byte[4];
        int count = readFully(lenField, 4);

        if (count != 4) {
            if (Krb5.DEBUG) {
                System.out.println(
                    ">>>DEBUG: TCPClient could not read length field");
            }
            return null;
        }

        int len = networkByteOrderToInt(lenField, 0, 4);
        if (Krb5.DEBUG) {
            System.out.println(
                ">>>DEBUG: TCPClient reading " + len + " bytes");
        }
        if (len <= 0) {
            if (Krb5.DEBUG) {
                System.out.println(
                    ">>>DEBUG: TCPClient zero or negative length field: "+len);
            }
            return null;
        }

        try {
            return IOUtils.readFully(in, len, true);
        } catch (IOException ioe) {
            if (Krb5.DEBUG) {
                System.out.println(
                    ">>>DEBUG: TCPClient could not read complete packet (" +
                    len + "/" + count + ")");
            }
            return null;
        }
    }

    @Override
    public void close() throws IOException {
        tcpSocket.close();
    }

    /**
     * Read requested number of bytes before returning.
     * @return The number of bytes actually read; -1 if none read
     */
    private int readFully(byte[] inBuf, int total) throws IOException {
        int count, pos = 0;

        while (total > 0) {
            count = in.read(inBuf, pos, total);

            if (count == -1) {
                return (pos == 0? -1 : pos);
            }
            pos += count;
            total -= count;
        }
        return pos;
    }

    /**
     * Returns the integer represented by 4 bytes in network byte order.
     */
    private static int networkByteOrderToInt(byte[] buf, int start,
        int count) {
        if (count > 4) {
            throw new IllegalArgumentException(
                "Cannot handle more than 4 bytes");
        }

        int answer = 0;

        for (int i = 0; i < count; i++) {
            answer <<= 8;
            answer |= ((int)buf[start+i] & 0xff);
        }
        return answer;
    }

    /**
     * Encodes an integer into 4 bytes in network byte order in the buffer
     * supplied.
     */
    private static void intToNetworkByteOrder(int num, byte[] buf,
        int start, int count) {
        if (count > 4) {
            throw new IllegalArgumentException(
                "Cannot handle more than 4 bytes");
        }

        for (int i = count-1; i >= 0; i--) {
            buf[start+i] = (byte)(num & 0xff);
            num >>>= 8;
        }
    }
}

class UDPClient extends NetClient {
    InetAddress iaddr;
    int iport;
    int bufSize = 65507;
    DatagramSocket dgSocket;
    DatagramPacket dgPacketIn;

    UDPClient(String hostname, int port, int timeout)
        throws UnknownHostException, SocketException {
        iaddr = InetAddress.getByName(hostname);
        iport = port;
        dgSocket = new DatagramSocket();
        dgSocket.setSoTimeout(timeout);
        dgSocket.connect(iaddr, iport);
    }

    @Override
    public void send(byte[] data) throws IOException {
        DatagramPacket dgPacketOut = new DatagramPacket(data, data.length,
                                                        iaddr, iport);
        dgSocket.send(dgPacketOut);
    }

    @Override
    public byte[] receive() throws IOException {
        byte[] ibuf = new byte[bufSize];
        dgPacketIn = new DatagramPacket(ibuf, ibuf.length);
        try {
            dgSocket.receive(dgPacketIn);
        }
        catch (SocketException e) {
            if (e instanceof PortUnreachableException) {
                throw e;
            }
            dgSocket.receive(dgPacketIn);
        }
        byte[] data = new byte[dgPacketIn.getLength()];
        System.arraycopy(dgPacketIn.getData(), 0, data, 0,
                         dgPacketIn.getLength());
        return data;
    }

    @Override
    public void close() {
        dgSocket.close();
    }
}