/*************************** GO-LICENSE-START********************************* * Copyright 2016 ThoughtWorks, Inc. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ************************GO-LICENSE-END***********************************/ package com.thoughtworks.go.websocket; import com.google.gson.*; import com.thoughtworks.go.domain.AgentRuntimeStatus; import com.thoughtworks.go.remote.AgentIdentifier; import com.thoughtworks.go.remote.work.Work; import com.thoughtworks.go.server.service.AgentBuildingInfo; import com.thoughtworks.go.server.service.AgentRuntimeInfo; import com.thoughtworks.go.server.service.ElasticAgentRuntimeInfo; import com.thoughtworks.go.util.StringUtil; import org.apache.commons.io.IOUtils; import java.io.*; import java.lang.reflect.Type; import java.nio.charset.StandardCharsets; import java.util.Base64; import java.util.zip.GZIPInputStream; import java.util.zip.GZIPOutputStream; import static com.thoughtworks.go.util.ExceptionUtils.bomb; public class MessageEncoding { private static Gson gson = new GsonBuilder().excludeFieldsWithoutExposeAnnotation().registerTypeAdapter(AgentRuntimeInfo.class, new AgentRuntimeInfoTypeAdapter()).create(); public static String encodeWork(Work work) { try { try (ByteArrayOutputStream binaryOutput = new ByteArrayOutputStream()) { try (ObjectOutputStream objectStream = new ObjectOutputStream(binaryOutput)) { objectStream.writeObject(work); } return Base64.getEncoder().encodeToString(binaryOutput.toByteArray()); } } catch (IOException e) { throw bomb(e); } } public static Work decodeWork(String data) { try { byte[] binary = Base64.getDecoder().decode(data.getBytes(StandardCharsets.UTF_8)); try (ObjectInputStream objectStream = new ObjectInputStream(new ByteArrayInputStream(binary))) { return (Work) objectStream.readObject(); } } catch (ClassNotFoundException | IOException e) { throw bomb(e); } } public static byte[] encodeMessage(Message msg) { String encode = gson.toJson(msg); try { try (ByteArrayOutputStream bytes = new ByteArrayOutputStream()) { try (GZIPOutputStream out = new GZIPOutputStream(bytes)) { out.write(encode.getBytes(StandardCharsets.UTF_8)); out.finish(); } return bytes.toByteArray(); } } catch (IOException e) { throw bomb(e); } } public static Message decodeMessage(InputStream input) { try { try (GZIPInputStream zipStream = new GZIPInputStream(input)) { String jsonStr = new String(IOUtils.toByteArray(zipStream), StandardCharsets.UTF_8); return gson.fromJson(jsonStr, Message.class); } } catch (IOException e) { throw bomb(e); } } public static String encodeData(Object obj) { return gson.toJson(obj); } public static <T> T decodeData(String data, Class<T> aClass) { return gson.fromJson(data, aClass); } // todo: Remove hand wrote deserialization after merging ElasticAgentRuntimeInfo class into AgentRuntimeInfo (@wpc) private static class AgentRuntimeInfoTypeAdapter implements JsonDeserializer<AgentRuntimeInfo> { @Override public AgentRuntimeInfo deserialize(JsonElement json, Type typeOfT, JsonDeserializationContext context) throws JsonParseException { JsonObject jsonObject = json.getAsJsonObject(); AgentIdentifier identifier = context.deserialize(jsonObject.get("identifier"), AgentIdentifier.class); AgentRuntimeStatus runtimeStatus = context.deserialize(jsonObject.get("runtimeStatus"), AgentRuntimeStatus.class); AgentBuildingInfo buildingInfo = context.deserialize(jsonObject.get("buildingInfo"), AgentBuildingInfo.class); String location = jsonObject.has("location") ? jsonObject.get("location").getAsString() : null; Long usableSpace = jsonObject.has("usableSpace") ? jsonObject.get("usableSpace").getAsLong() : null; String operatingSystemName = jsonObject.has("operatingSystemName") ? jsonObject.get("operatingSystemName").getAsString() : null; String cookie = jsonObject.has("cookie") ? jsonObject.get("cookie").getAsString() : null; boolean supportsBuildCommandProtocol = jsonObject.has("supportsBuildCommandProtocol") && jsonObject.get("supportsBuildCommandProtocol").getAsBoolean(); String elasticPluginId = jsonObject.has("elasticPluginId") ? jsonObject.get("elasticPluginId").getAsString() : null; String elasticAgentId = jsonObject.has("elasticAgentId") ? jsonObject.get("elasticAgentId").getAsString() : null; AgentRuntimeInfo info; if (elasticPluginId == null || StringUtil.isBlank(elasticPluginId)) { info = new AgentRuntimeInfo(identifier, runtimeStatus, location, cookie, supportsBuildCommandProtocol); } else { info = new ElasticAgentRuntimeInfo(identifier, runtimeStatus, location, cookie, elasticAgentId, elasticPluginId); } info.setUsableSpace(usableSpace); info.setOperatingSystem(operatingSystemName); info.setSupportsBuildCommandProtocol(supportsBuildCommandProtocol); info.setBuildingInfo(buildingInfo); return info; } } }